diff --git a/.claude/skills/band-code-review/angle-groups.md b/.claude/skills/band-code-review/angle-groups.md index 1bc52381a..72664715b 100644 --- a/.claude/skills/band-code-review/angle-groups.md +++ b/.claude/skills/band-code-review/angle-groups.md @@ -37,15 +37,34 @@ and consequences differ. contents, secrets). **D — Test Quality** -Test code gets every angle above, more strictly, plus: -- Apply `band-python-test-conventions` "Coverage and realistic flows" to the - changed behavior and relevant existing tests, even when no test file changed. - Map scenarios to meaningful assertions, identify gaps, and name the concrete - regression each proposed test would catch. Check complete complex flows at - a boundary that preserves their real interactions; passing tests or high - line coverage alone are not sufficient evidence. -- A skipped, disabled, or silently-ignored failing test is always a blocking - issue (`band-python-test-conventions`). +Test code belongs here alone. Read the whole `band-python-test-conventions` +skill, not only "Coverage is branches, not lines". Apply A–C more strictly +than on production code (SSoT, intent-named helpers, no mixed-concern files, +no magic wire strings). Each of the following is an `issue:`, not a nit: + +- Coverage of complementary branches, not line counts. Map each changed + `match`/`if`/timeout/OS decision to the test that proves each reachable + outcome. Proving fail-closed does not prove the success side; a + `skipif(os.name)` is a missing branch on the skipped OS. Name the + regression each proposed test would catch. Passing tests or high line + coverage alone are not evidence. +- Assertions of plumbing or assumptions: occupancy arithmetic, private + fields, mock call counts, event-type tallies as a proxy, empty collections + the test does not own, or tautologies. The assertion must be a named + projection of the decision. +- A test function that branches on scenario/kind (`if scenario != ...`) or + re-types a wire literal the peer or product already names. Parametrize + `(scenario, expected)` so each row *is* the claim. +- The interaction under test is patched away (`unittest.mock.patch` of + private names, stdlib spawn, or timeout constants) instead of a public + seam or a real peer. `AsyncMock` of the unit under test is the same smell. +- Wall-clock sleep or waiting out a grace period used as proof. + `asyncio.timeout` is a hung-test guard only; a public timeout of `0` is + "check now", and it still needs the complementary "finished before wait" + sequence. +- A skipped, disabled, or silently-ignored failing test. +- A Hypothesis model or kit file that looks like a collected test (a `Test*` + class, or a missing role docstring, in a non-`test_*.py` module). **E — Edge-Case Flow Analysis** Owns failures that need a *sequence* of state transitions or interactions. A diff --git a/.claude/skills/band-here-is-a-review/SKILL.md b/.claude/skills/band-here-is-a-review/SKILL.md index f427e4cb1..5e23385e0 100644 --- a/.claude/skills/band-here-is-a-review/SKILL.md +++ b/.claude/skills/band-here-is-a-review/SKILL.md @@ -109,7 +109,9 @@ Only CONFIRMED and PARTIALLY-RIGHT findings that **matter** proceed to fixing. "Matters" means: user-visible breakage, data corruption, security exposure, a real maintenance trap, or a violated project invariant or rule (`AGENTS.md` and the guideline skills `band-python-conventions` and `band-python-test-conventions` — e.g. duplicated logic, -magic literals, narrating comments). Filter out: +magic literals, narrating comments, tests that assert plumbing or assumptions, +patch internals, wait out wall-clock grace, branch on scenario inside the test, +or treat line coverage as branch coverage). Filter out: - Pedantry and taste conflicts with the project's established conventions (the codebase's existing idiom wins over the reviewer's preference). diff --git a/.claude/skills/band-plan/plan-template.md b/.claude/skills/band-plan/plan-template.md index 599249381..03c386044 100644 --- a/.claude/skills/band-plan/plan-template.md +++ b/.claude/skills/band-plan/plan-template.md @@ -52,12 +52,20 @@ Test rules: - Test behavior that can actually break, not defaults or implementation trivia (see the `band-python-test-conventions` skill). A test that would still pass under a symptom-level patch doesn't cover the root cause. + Cover complementary branches of each owned decision, not line execution; + assert a named projection of the outcome, not internals or unclaimed + empties. Parametrize `(scenario, expected)` — do not `if` on scenario + inside the test. - When a material scenario needs end-to-end coverage to prove the behavior, reuse or extend the `tests/e2e/baseline` toolkit (see the `band-write-e2e-test` skill) instead of building new plumbing. Otherwise use the lowest boundary that preserves the behavior being tested. -- A missing helper, fixture, or local peer is its own step, before the step - that uses it — never an empty mock in place of a real boundary. +- A missing helper, fixture, constructor seam, or local peer is its own + step, before the step that uses it — never `unittest.mock.patch` of a + private name, stdlib spawn, or timeout constant, and never an empty mock + in place of a real boundary. Timeouts are proven with a public `0` ("check + now") plus the complementary "finished before wait" sequence, not by + sleeping the grace period. Scope rules: diff --git a/.claude/skills/band-python-conventions/SKILL.md b/.claude/skills/band-python-conventions/SKILL.md index 88dfe5e5f..bf209eb4a 100644 --- a/.claude/skills/band-python-conventions/SKILL.md +++ b/.claude/skills/band-python-conventions/SKILL.md @@ -1,6 +1,6 @@ --- name: band-python-conventions -description: Band Python SDK conventions for Python source, comments, docstrings, and error handling. Read before editing Python in this repository. +description: Band Python SDK conventions for Python source, comments, docstrings, and error handling. Read before editing Python in this repository. Test files also follow band-python-test-conventions; this skill still owns SSoT, intent-named helpers, and public test seams. paths: - "**/*.py" --- @@ -44,7 +44,12 @@ paths: - Use `%s` placeholders in log messages for lazy evaluation - Use Pydantic v2 for data models; use `model_dump()` not `dict()` - Target Python 3.11+; use `list[str]` not `List[str]`, `str | None` not `Optional[str]` -- Use async/await everywhere in async codebases; use `AsyncMock` for testing async methods +- Use async/await everywhere in async codebases. `AsyncMock` is a test double + at an external boundary, not a stand-in for the unit whose interaction a + test claims to prove. Public constructor seams (a `create_process` factory, + an instance timeout of `0` meaning "check now") beat `unittest.mock.patch` + of private names or stdlib spawn. Tests follow + `band-python-test-conventions`. - Catch `pydantic.ValidationError` separately from generic `Exception` - Use `raise ValueError(...)` for missing required config, not `logger.error()` + `sys.exit()` - Never put issue-tracker references in code — no Linear issue IDs (e.g. `INT-123`), Linear URLs, or ticket numbers in comments, docstrings, or strings. Explain the *why* in plain terms instead. (Branch names, commit messages, and PR descriptions may reference issues.) @@ -62,10 +67,12 @@ paths: - Prefer a single source of truth for a value or closed vocabulary consumed in more than one place: give it one definition — a constant, a `StrEnum`, or a small helper — that every site references, rather than re-typing the same magic literal in a - producer and the consumer that reads it (a typo then fails silently). Keep genuinely - distinct vocabularies separate, though — don't merge two sets that only happen to - share some values today (e.g. the ACP `ChunkType` a chunk carries vs. the platform - `message_type` an event is posted under). + producer and the consumer that reads it (a typo then fails silently). This includes + tests: a wire status, prompt text, or error message typed in a peer and again in an + assertion is two copies. Keep genuinely distinct vocabularies separate, though — + don't merge two sets that only happen to share some values today (e.g. the ACP + `ChunkType` a chunk carries vs. the platform `message_type` an event is posted + under). - Before writing new logic, search the codebase for an existing helper that does or almost does it; extending it beats a parallel implementation. - Use idiomatic constructs (`match`, decorators, context managers, recursion) diff --git a/.claude/skills/band-python-test-conventions/SKILL.md b/.claude/skills/band-python-test-conventions/SKILL.md index 7287c9350..aa4bdcf88 100644 --- a/.claude/skills/band-python-test-conventions/SKILL.md +++ b/.claude/skills/band-python-test-conventions/SKILL.md @@ -1,6 +1,6 @@ --- name: band-python-test-conventions -description: Band Python SDK test conventions and coverage assessment. Use when writing or reviewing tests, assessing whether changed behavior is adequately covered, or checking that tests validate meaningful outcomes and realistic complex flows. +description: Band Python SDK test conventions and coverage assessment. Use when writing or reviewing tests, assessing whether changed behavior is adequately covered, or checking that tests validate meaningful outcomes and realistic complex flows. Also use when a test asserts internals, patches private names, waits on wall-clock grace, branches on scenario inside the test, or reports line coverage instead of complementary branches. paths: - "tests/**/*.py" - "src/**/*.py" @@ -8,71 +8,165 @@ paths: # Test Conventions -## Coverage and realistic flows +`band-python-conventions` still applies: single source of truth, intent-named +helpers, and no mixed-concern files. This skill is the extra bar for tests. + +## Coverage is branches, not lines Assess the relevant existing tests whenever behavior changes or coverage is reviewed, even if the diff adds no tests. A green suite, a test count, or a line-coverage percentage alone does not establish adequate coverage. -- Map user-visible scenarios and important invariants to the tests and - assertions that prove them. Identify missing or weak coverage, prioritizing - consequential behavior, changed branches, and regressions over a numeric target. -- For each test, name the plausible bug it would catch. Assert the observable - outcome and relevant side effects; a test that only checks that a mock was - called or that nothing raised may leave the real behavior unverified. -- Exercise complex flows as realistic sequences, not just isolated steps. - Where relevant, cover follow-up turns, failure and recovery, cancellation, - restart/rehydration, concurrent rooms or agents, permissions, and cleanup. - Select scenarios the actual feature can reach; do not add every case to - every change. -- Drive the real entry point and interacting components at the lowest level - that preserves the behavior under test. Use realistic peers or fakes at - external boundaries; do not mock away the interaction whose correctness - the test claims to prove. Use the baseline toolkit for needed live E2E - coverage, not as a mandatory layer for every test. -- Extend useful existing tests before duplicating them. Add or strengthen - tests for meaningful in-scope gaps; report remaining gaps and unverified - scenarios explicitly. Follow `AGENTS.md` for synchronization and validation - scope, without silently expanding the task or launching unrelated suites. - -## Test implementation +Map each production decision the change owns — a `match`/`case`, an `if` that +chooses an outcome, an OS spawn path, a timeout pending-vs-done wait — onto +the tests that prove it. For every such decision: + +1. Name the reachable outcomes (succeed, fail-closed, ignore, retry, OS-A vs + OS-B). +2. Name the test that proves each as an observable sequence, not as a line + that executed. +3. Treat complementary pairs as separate scenarios. Proving one side does not + prove the other: backlog live-at-limit vs overflow; a callback that + finishes before the wait vs one still running; settlement before the + matching result vs result before settlement; POSIX process-group vs + Windows job. +4. Report remaining unexamined outcomes explicitly. A Hypothesis model covers + generated sequences of its Step alphabet, not branches those steps never + take. + +Prioritize consequential behavior, changed branches, and regressions over a +numeric target. Do not add every case to every change; select scenarios the +actual feature can reach. + +For each test, name the plausible bug it would catch. Assert the observable +outcome and relevant side effects; a test that only checks that a mock was +called or that nothing raised may leave the real behavior unverified. + +Exercise complex flows as realistic sequences, not just isolated steps. +Where relevant, cover follow-up turns, failure and recovery, cancellation, +restart/rehydration, concurrent rooms or agents, permissions, and cleanup. + +Extend useful existing tests before duplicating them. Follow `AGENTS.md` for +synchronization and validation scope, without silently expanding the task or +launching unrelated suites. + +## Assertions show intent, not plumbing or assumptions + +Tests must be declarative and intent-revealing, not a transcript of the +implementation. Assert on a readable projection of the observable outcome — +the thing the test is actually about — never on raw internals or on a side +effect that merely implies the real answer. + +- `assert reply.outline == ["tool_call (permission)", "message", ...]` over a + hand-rolled comprehension pulling `message_type` out of each event dict. +- `assert record.levelno == logging.DEBUG` over inferring a log level + indirectly from whether two separate capture windows came back empty. +- A named helper (`owned_rpc_process`, `outline`) over poking `_proc` or + counting `AGENT_END` as a proxy for "the turn finished twice". +- Overload of `EVENT_BACKLOG_LIMIT + 1` over occupancy arithmetic + (`limit - 2`) that restates the queue's internals. + +If writing the assertion requires re-deriving *how* the code decided +something, the test is checking the wrong thing — assert the decision +itself. + +Do not assert an assumption the test does not own. Drop empty-collection +checks on a mapping test that is about session keys, tautologies such as +Windows `(None, None)`, and "observer stayed empty" when the claim is +fail-closed-and-restarted. If the assertion would still pass after inverting +the behavior on an unclaimed axis, it is noise, not coverage. + +Test what really matters — behavior that can break. Don't write tests that +restate definitions (asserting dataclass defaults equal themselves, echoing a +constant) or otherwise cannot fail for a real reason. + +Decide what correct looks like before running the code, then assert that; +never bend an assertion to match whatever the code currently produces. + +## One concern per test, no in-test branching + +Each `test_*.py` holds exactly one `class Test...` named for the behavior it +covers — no module-level test functions and no helper classes beside it. +When a suite outgrows one file, make it a package: `conftest.py` for +fixtures, a plain helper module (`kit.py`, a Hypothesis `lifecycle.py`) for +shared doubles, and one test file per concern. + +A helper module in that package is not a collected test. Its module +docstring names the role (kit, peer, lifecycle model). It does not define a +`Test*` class. The collected test that *plays* a Hypothesis model lives in +`test_lifecycle_model.py` (see `tests/platform/linkclose/test_lifecycle_model.py`). + +One test function, one claim. Parametrize `(scenario, error)` so each row's +assertion *is* the claim. A long test that dispatches with `if scenario != +"aborted"` or `if os.name == "nt": return` is several tests hiding in one +function; split them. + +Wire statuses, prompt texts, and error messages belong to one named constant +or enum that the peer, the product, and the assertion all reference. A +hard-coded `"complete"` / `"Approve"` in the test is a second copy of a +closed vocabulary. + +## Drive the real boundary; inject, don't patch + +Drive the real entry point and interacting components at the lowest level +that preserves the behavior under test. Use realistic peers or fakes at +external boundaries; do not mock away the interaction whose correctness the +test claims to prove. + +- A constructor seam (`create_process=`) or a public instance timeout of `0` + ("check now / no grace") is the test hook. `unittest.mock.patch` of + `asyncio.create_subprocess_exec`, a private `_close_process`, or a + `_CALLBACK_EXIT_TIMEOUT_S` constant is not. +- `AsyncMock` is a double at an external boundary (a client the function + *takes*), not a stand-in for the unit under test. A receipt-table test may + mock the RPC client; it then does not count as coverage of the client's + prompt path. +- Use the baseline toolkit for needed live E2E coverage, not as a mandatory + layer for every test. + +## Time is a hung-test guard, not proof + +Ordering, concurrency, and timeout behavior is proven with explicit +synchronization and a controlled clock (`AGENTS.md` "Test Synchronization"): +events, gates, barriers, or observable state transitions. Never use elapsed +wall time, `asyncio.sleep` of a grace period, or "it finished within a +duration" as proof. + +- `async with asyncio.timeout(N)` around a test is a hung-test guard. It + must not be the mechanism that establishes the outcome. +- To prove a timeout fail-closed path, set the public timeout to `0` and + assert the failure. Do not wait out the production grace period. +- `timeout = 0` proves "still pending at the wait". It does not prove the + complementary branch: work that finishes before the wait succeeds. Cover + that with a sequence that completes, not with a longer sleep. + +Real subprocess and network readiness stay independent of a controlled +clock. Cleanup must account for resources created during failed or cancelled +startup and await their confirmed release. + +## Tests run on Ubuntu and Windows CI + +Test across OSes; don't skip one. Build host paths with `tmp_path` or +`tests.paths.host_absolute_path`, never a hard-coded `"/opt/..."` (no drive, +so relative on Windows); check absoluteness with +`band.workspaces.is_host_absolute`. + +Where behavior differs per OS, write one parametrize table that runs +everywhere, each row naming the `os.name` values that accept it. +`skipif(os.name ...)` is only for a boundary that cannot run on that OS, +never for a path, separator, or "this assertion is awkward on Windows". A +skip is a missing branch on the skipped OS. + +## Paths, bugs, skips - In tests, never derive repository-anchored paths with per-file `Path(__file__).parents[N]` arithmetic — import the anchors from `tests/paths.py` (`REPO_ROOT`, `SRC_ROOT`, `EXAMPLES_ROOT`, `KIT_DIR`, `ENV_TEST_FILE`). Only genuinely package-relative paths (a fixture file next to its test) stay relative to their own `__file__`. -- Test what really matters — behavior that can break. Don't write tests that - restate definitions (asserting dataclass defaults equal themselves, echoing a - constant) or otherwise cannot fail for a real reason; they add maintenance - cost without protection. -- **Tests must be declarative and intent-revealing, not a transcript of the - implementation.** Assert on a readable projection of the observable outcome - — the thing the test is actually about — never on raw internals or on a - side effect that merely implies the real answer. Concretely: - - `assert reply.outline == ["tool_call (permission)", "message", ...]` over - a hand-rolled comprehension pulling `message_type` out of each event dict. - - `assert record.levelno == logging.DEBUG` over inferring a log level - indirectly from whether two separate capture windows came back empty. - If writing the assertion requires re-deriving *how* the code decided - something, the test is checking the wrong thing — assert the decision - itself. -- **Tests run on Ubuntu and Windows CI — test across OSes, don't skip one.** - Build host paths with `tmp_path` or `tests.paths.host_absolute_path`, never a - hard-coded `"/opt/..."` (no drive, so relative on Windows); check absoluteness - with `band.workspaces.is_host_absolute`. Where behavior differs per OS, write - one parametrize table that runs everywhere, each row naming the `os.name` - values that accept it. `skipif(os.name ...)` is only for a boundary that - cannot run on that OS, never for a path or separator assumption. -- Decide what correct looks like before running the code, then assert that; - never bend an assertion to match whatever the code currently produces. - A skipped, disabled, or silently ignored failing test is always a blocking issue. - A fix for a confirmed bug comes with a test that fails before the fix and passes after it. - For a stateful lifecycle component, add a Hypothesis model-based test over - generated step sequences (see `tests/platform/linkclose/test_lifecycle_model.py`). -- Ordering, concurrency and timeout behavior is proven with explicit - synchronization and a controlled clock: `AGENTS.md` "Test Synchronization". -- File layout (one test class per file, packages for larger suites): - `AGENTS.md` "Testing Structure". + generated step sequences. diff --git a/.claude/skills/band-write-e2e-test/SKILL.md b/.claude/skills/band-write-e2e-test/SKILL.md index ec98bb44f..54cf6a00b 100644 --- a/.claude/skills/band-write-e2e-test/SKILL.md +++ b/.claude/skills/band-write-e2e-test/SKILL.md @@ -79,9 +79,9 @@ reuse, target location, precedents, and any coverage overlap). ## Phase 2 — Plan the test Apply [test conventions](../band-python-test-conventions/SKILL.md) -"Coverage and realistic flows" when mapping existing coverage and selecting -the scenario. Show which meaningful gap this test closes and which complete -interaction it proves. +"Coverage is branches, not lines" when mapping existing coverage and selecting +the scenario. Show which complementary-branch gap this test closes and which +complete interaction it proves. Write a concrete plan (not code yet). It must answer: diff --git a/AGENTS.md b/AGENTS.md index ffb1ee70d..55a5541b4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -80,7 +80,7 @@ explicit user approval. | Skill | Use when | |---|---| | [band-python-conventions](.claude/skills/band-python-conventions/SKILL.md) | Python source, comments, docstrings, and error handling | -| [band-python-test-conventions](.claude/skills/band-python-test-conventions/SKILL.md) | writing/reviewing tests or assessing meaningful coverage and realistic flows | +| [band-python-test-conventions](.claude/skills/band-python-test-conventions/SKILL.md) | writing/reviewing tests; complementary branches not line counts; intent assertions; seams not patches; no wall-clock proof | | Skill | Use when | |---|---| @@ -447,7 +447,10 @@ install` it and exercise the real call in this repo's venv. For Python changes, before selecting checks, re-read your own diff once against the `band-python-conventions` skill — including code you just wrote this session, -not only code you started from. The two rules that get skipped under time pressure: +not only code you started from. Test diffs also re-read +`band-python-test-conventions` (complementary branches, intent assertions, +seams not patches, no wall-clock proof). The two rules that get skipped under +time pressure: - **Single source of truth**: a literal, magic string, or multi-line block re-typed in more than one place (a second copy you just wrote counts) instead diff --git a/CHANGELOG.md b/CHANGELOG.md index edac55b1f..8a0117a0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,15 @@ The format is based on [Conventional Commits](https://www.conventionalcommits.or and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). This changelog is automatically generated by [Release Please](https://github.com/googleapis/release-please). +## Unreleased + +### Bug Fixes + +* Wait for Windows Job members' kernel-confirmed exit before releasing OMP or Codex subprocess ownership, preserving process handles across incomplete cleanup and preventing late child admission during termination. +* Retire native OMP sessions on an observed root exit, preserve usable sessions after post-admission command errors, and keep callbacks admitted during settlement within the original exit deadline. +* Keep native OMP stderr payloads out of SDK debug logs. +* Preserve false and empty native OMP UI answers without adding a cancellation flag, and ignore whitespace-only connection instructions. + ## [6.0.1](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v6.0.0...band-sdk-v6.0.1) (2026-10-20) diff --git a/docs/acp.md b/docs/acp.md index 5ee33cc95..3b120f95f 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -281,3 +281,41 @@ cleanup. Explicit caller cancellation and provider timeouts remain effective. credential-redacted; values under credential-named fields are redacted in full. - The ACP Python client drains queued update handlers before surfacing a prompt error, so a slow client handler can still delay its own `prompt()`. + +## Native OMP transport + +`band.integrations.omprpc.OmpRpcClient` speaks OMP's native RPC v2 protocol, +including structured ask questions and answers. It is separate from the ACP +transport; adding it does not deprecate or remove the public `OmpACPAdapter`. + +Event callbacks run serially with a bounded backlog. Exhausting that backlog +fails the owning prompt with `OmpProcessError` and starts cleanup rather than +silently dropping events and reporting success. The wire reader does not wait +for callback capacity, so callbacks can still issue correlated commands and +answer native UI requests while delivery is held. + +An observed root-process exit also retires the transport when descendants still +hold stdout open. Command errors, including an error reported after prompt +admission, raise `OmpRpcError` without retiring an otherwise usable process. +Host-tool and UI callbacks admitted while a settled turn drains share its +original callback-exit deadline; an unfinished callback fails the turn and starts +cleanup. SDK debug logs report stderr byte counts, never the child's raw stderr. + +Explicit MCP transports use `band.integrations.ompprocess.OmpMcpTransport`. +Native session metadata uses `band.converters.omp.OmpMetadataKey.ROOM` and +`OmpMetadataKey.SESSION`. + +A `create_process` factory customizes creation without bypassing containment. +On Windows it launches the existing helper; the client still owns the Job +containing the native process and its descendants, including during cancelled +or incomplete startup. +To hold or observe the owned launch, wrap `spawn_owned_process` and delegate to +its bound original (or override it and call `super()`). Do not pass that wrapper +as `create_process`: that factory runs below the ownership boundary. The owned +method retains the process before returning, so wrapper failure still reaps it. + +On POSIX, process-group capture rechecks a reaped leader's identity after +enumerating members, before retaining that group. A reused leader PID does not +transfer custody to an unrelated session. Platforms without process-group support +use their existing Job and environment-marker ownership instead of probing that +PID as a group identity. diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index 41d6e5ce8..adbb58605 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -35,12 +35,17 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra close errors and close timeouts retain the claim for retry. A timed-out or cancelled caller does not abandon subprocess cleanup; retry room cleanup to finish releasing ownership. Shutdown attempts every room and reports failures. + On POSIX, close sends stdin EOF and gives the session a bounded chance to + exit before reaping remaining descendants. Windows shutdown does not promise + session-visible stdin EOF. On Windows, the client owns the helper and its process-tree job from startup through cleanup. It starts the interpreter image directly so a virtual-environment launcher cannot hide another helper process. Cleanup confirms the helper has - exited before terminating and draining the job; an empty job alone does not - establish helper exit. - Failed cleanup retains that same owner and job for retry. + exited, prevents new job members from being admitted, and pins the current + members' process identities before terminating the job. It waits for both empty + job accounting and kernel-confirmed member exit; an empty job alone does not + establish process exit. Failed cleanup retains that same owner, job and process + handles for retry. - **Leaving is final.** Leaving a room stops it admitting work at once. A leave that fails or runs over keeps the process, thread and workspace claim, and the room refuses work until a retried leave succeeds; a later rejoin diff --git a/pyproject.toml b/pyproject.toml index 1272669c8..ff58e1638 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,9 @@ acp = [ "starlette>=0.40.0", "uvicorn>=0.29", ] +omp = [ + "psutil>=6", +] slack = [ "slack-sdk>=3.27.0", # AsyncWebClient, signature verifier, Socket Mode "starlette>=0.40.0", # HTTP server framework @@ -206,6 +209,7 @@ dev = [ "pytest-rerunfailures>=14.0.0", "tenacity>=8.0.0", "pytest-timeout>=2.4.0", + "psutil>=6", # Native OMP process-tree cleanup tests "packaging>=23.0", # PEP 440 pre-release detection for the CI coverage-gate version guard "band-testing-python==0.1.4", "httpx>=0.24.0", # Already in main deps, for mocking diff --git a/src/band/converters/__init__.py b/src/band/converters/__init__.py index 431fb6876..bde4b094d 100644 --- a/src/band/converters/__init__.py +++ b/src/band/converters/__init__.py @@ -98,6 +98,12 @@ from band.converters.nooa import ( NooaTranscript as NooaTranscript, ) + from band.converters.omp import ( + OmpHistoryConverter as OmpHistoryConverter, + ) + from band.converters.omp import ( + OmpSessionState as OmpSessionState, + ) from band.converters.opencode import ( OpencodeHistoryConverter as OpencodeHistoryConverter, ) @@ -139,6 +145,7 @@ gemini=["GeminiHistoryConverter", "GeminiMessages"], google_adk=["GoogleADKHistoryConverter", "GoogleADKMessages"], opencode=["OpencodeHistoryConverter"], + omp=["OmpHistoryConverter", "OmpSessionState"], nooa=["NooaHistoryConverter", "NooaTranscript"], strands=["StrandsHistoryConverter", "StrandsMessages"], ) diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py new file mode 100644 index 000000000..e83f643c6 --- /dev/null +++ b/src/band/converters/omp.py @@ -0,0 +1,43 @@ +"""Native OMP resume candidates and read-only room transcript replay.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, Field + +from band.converters.helpers import build_replay_messages +from band.core.protocols import HistoryConverter +from band.core.types import MessageType + + +class OmpMetadataKey(StrEnum): + """Task-event metadata keys the native adapter writes and this converter reads.""" + + ROOM = "omp_room_id" + SESSION = "omp_session_file" + + +class OmpSessionState(BaseModel): + """The native session file each room may resume, plus the room's transcript.""" + + room_to_session: dict[str, str] = Field(default_factory=dict) + replay_messages: list[str] = Field(default_factory=list) + + +class OmpHistoryConverter(HistoryConverter[OmpSessionState]): + """Extract per-room native session files from task-event metadata.""" + + def convert(self, raw: list[dict[str, Any]]) -> OmpSessionState: + sessions: dict[str, str] = {} + for message in raw: + match message.get("message_type"), message.get("metadata"): + case MessageType.TASK, { + OmpMetadataKey.ROOM: str(room), + OmpMetadataKey.SESSION: str(session), + } if session: + sessions[room] = session + return OmpSessionState( + room_to_session=sessions, replay_messages=build_replay_messages(raw) + ) diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index c87961e5a..85c9a92bc 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -7,10 +7,7 @@ import logging import os import shutil -import signal from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import suppress -from pathlib import Path from typing import Any from band.core.exceptions import BandConnectionError @@ -21,8 +18,7 @@ WindowsProcess, await_returncode, close_stdin, - find_executable, - process_image, + kill_process_tree, ) logger = logging.getLogger(__name__) @@ -30,55 +26,10 @@ # Long enough for the child watcher to record an exit, short enough that a # held stdio pipe cannot stall room cleanup. _EXIT_OBSERVE_S = 2.0 +# How long a closed stdin gets to end the whole Codex session before it is killed. +_SESSION_EXIT_S = 1.0 +_SESSION_POLL_S = 0.02 _STDOUT_LIMIT = 16 * 1024 * 1024 -HELPER_PATH = Path(__file__).resolve().with_name("winmain.py") - - -def _absolute_command(command: Sequence[str]) -> tuple[str, ...]: - """Resolve the executable in this process, not in the room directory. - - The helper's current directory is the room. ``CreateProcessW`` searches - that directory when the application name is relative, so a room-local - binary would run in place of the one this process would have started. - """ - head, *rest = command - if Path(head).is_absolute(): - return (head, *rest) - if os.name == "nt" and not os.path.dirname(head): - found = find_executable(head) - else: - found = shutil.which(head) - if found is None and Path(head).exists(): - found = str(Path(head).resolve()) - if found is None: - raise FileNotFoundError(head) - return (str(Path(found).resolve()), *rest) - - -def _session_pid_reused(process: asyncio.subprocess.Process) -> bool: - """True when ``process.pid`` now names some other live process. - - After the leader is reaped, its session id stays reserved only while a - member of that session still exists. Once the id is free it can be given - to a new session, and ``killpg`` would signal that session. - """ - if process.returncode is None: - return False - try: - os.kill(process.pid, 0) - except ProcessLookupError: - return False - except OSError: - logger.warning( - "Codex process %s pid is not signalable; leaving its session", - process.pid, - ) - return True - logger.warning( - "Codex process %s pid was reused; not signaling that session", - process.pid, - ) - return True def _missing_cli_message(binary: str) -> str: @@ -98,18 +49,28 @@ async def _wait_readers(readers: Sequence[asyncio.Task[None]]) -> None: logger.warning("Codex stdio readers still pending after tree kill") -async def kill_process_tree(process: asyncio.subprocess.Process) -> None: - """Force-stop a Unix process and everything in its session. +async def _await_session_exit( + process: asyncio.subprocess.Process, timeout: float +) -> bool: + """True once the child and every other process in its session are gone. - The codex launcher is a shim whose real CLI is its child. On Unix the - child leads a session, so ``killpg`` still reaches that group after the - leader has exited. + The session id stays reserved while any member lives, so probing the + group is exact; the launcher alone exiting is not enough, since its + real CLI child would keep the room's working directory. """ - if os.name != "nt": - if not _session_pid_reused(process): - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - return + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while True: + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return True + except PermissionError: + pass + remaining = deadline - loop.time() + if remaining <= 0: + return False + await asyncio.sleep(min(_SESSION_POLL_S, remaining)) class CodexStdioClient(BaseJsonRpcClient): @@ -177,8 +138,8 @@ async def _spawn(self) -> asyncio.subprocess.Process: raise BandConnectionError("Codex process is already starting") self._windows_process = owner = WindowsProcess() try: - return await owner.start( - self._windows_command, + return await owner.start_command( + self.command, cwd=self.cwd, env=self.env, limit=_STDOUT_LIMIT, @@ -198,18 +159,6 @@ async def _spawn(self) -> asyncio.subprocess.Process: limit=_STDOUT_LIMIT, ) - def _windows_command(self, job_name: str, status_handle: int) -> list[str]: - return [ - str(process_image()), - "-I", - "-S", - str(HELPER_PATH), - job_name, - str(status_handle), - "--", - *_absolute_command(self.command), - ] - def _spawn_connection_error(self, failure: SpawnFailure) -> BandConnectionError: if failure.winerror in MISSING_EXECUTABLE: return BandConnectionError(_missing_cli_message(self.command[0])) @@ -258,8 +207,11 @@ async def _close_process(self) -> None: await self._windows_process.close(_EXIT_OBSERVE_S) self._windows_process = None elif (process := self._proc) is not None: - await kill_process_tree(process) + # Codex and its launcher exit on their own once stdin closes; the + # tree kill is the fallback for a session that does not. await close_stdin(process) + if not await _await_session_exit(process, _SESSION_EXIT_S): + await kill_process_tree(process) if not await await_returncode(process, _EXIT_OBSERVE_S): logger.warning( "Codex process %s still alive after tree kill", process.pid diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index a7306ba9c..e85b4c62f 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -3,6 +3,11 @@ The parent retains the job handle. The isolated helper creates the real command in that job before it can run, then waits. Close terminates the job so descendants die even after their launcher exits. + +This is a single kill-on-close job with no breakaway and no nested jobs. +Children inherit membership at create; they must not inherit the handle +(or kill-on-close never fires). Nested jobs or BREAKAWAY_OK belong here +only when close still leaves members alive. """ from __future__ import annotations @@ -12,32 +17,42 @@ import functools import logging import os +import shutil +import signal import subprocess import sys import uuid -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, suppress from ctypes import wintypes from pathlib import Path from typing import Any, NoReturn +CreateProcess = Callable[..., Awaitable[asyncio.subprocess.Process]] + logger = logging.getLogger(__name__) # CreateProcess failed because the executable path does not exist. MISSING_EXECUTABLE = frozenset({2, 3}) +HELPER_PATH = Path(__file__).resolve().with_name("winmain.py") _READY = "ok" _STATUS_ERROR = "error" _SPAWN_ACK_S = 30.0 _JOB_POLL_S = 0.05 +_JOB_LIST_QUERIES = 2 _ERROR_ALREADY_EXISTS = 183 _ERROR_BROKEN_PIPE = 109 +_ERROR_INVALID_PARAMETER = 87 +_ERROR_MORE_DATA = 234 _STATUS_READ_SIZE = 256 _PATH_BUFFER_SIZE = 260 _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 +_JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x00000008 _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9 _JOB_OBJECT_BASIC_ACCOUNTING_INFORMATION = 1 +_JOB_OBJECT_BASIC_PROCESS_ID_LIST = 3 _JOB_OBJECT_ASSIGN_PROCESS = 0x0001 _CREATE_SUSPENDED = 0x00000004 @@ -46,6 +61,9 @@ _HANDLE_FLAG_INHERIT = 0x00000001 _PROC_THREAD_ATTRIBUTE_HANDLE_LIST = 0x00020002 _PROC_THREAD_ATTRIBUTE_JOB_LIST = 0x0002000D +_PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 +_SYNCHRONIZE = 0x00100000 +_WAIT_OBJECT_0 = 0 _INFINITE = 0xFFFFFFFF _WAIT_FAILED = 0xFFFFFFFF @@ -116,6 +134,14 @@ class BasicAccountingInformation(ctypes.Structure): ) +class ProcessIdList(ctypes.Structure): + _fields_ = ( + ("NumberOfAssignedProcesses", wintypes.DWORD), + ("NumberOfProcessIdsInList", wintypes.DWORD), + ("ProcessIdList", ctypes.c_size_t * 1), + ) + + class StartupInfo(ctypes.Structure): _fields_ = ( ("cb", wintypes.DWORD), @@ -201,6 +227,14 @@ def _kernel32() -> Any: library.CloseHandle.restype = wintypes.BOOL library.OpenJobObjectW.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.LPCWSTR] library.OpenJobObjectW.restype = handle + library.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + library.OpenProcess.restype = handle + library.IsProcessInJob.argtypes = [ + handle, + handle, + ctypes.POINTER(wintypes.BOOL), + ] + library.IsProcessInJob.restype = wintypes.BOOL library.CreateProcessW.argtypes = [ wintypes.LPCWSTR, ctypes.POINTER(ctypes.c_wchar), @@ -344,6 +378,37 @@ def find_executable(name: str) -> str | None: capacity = length + 1 +def absolute_command(command: Sequence[str]) -> tuple[str, ...]: + """Resolve the executable in this process, not in the command's directory.""" + head, *rest = command + if Path(head).is_absolute(): + return (head, *rest) + if os.name == "nt" and not os.path.dirname(head): + found = find_executable(head) + else: + found = shutil.which(head) + if found is None and Path(head).exists(): + found = str(Path(head).resolve()) + if found is None: + raise FileNotFoundError(head) + return (str(Path(found).resolve()), *rest) + + +def command_in_job( + command: Sequence[str], job_name: str, status_handle: int +) -> list[str]: + return [ + str(process_image()), + "-I", + "-S", + str(HELPER_PATH), + job_name, + str(status_handle), + "--", + *absolute_command(command), + ] + + def _winerror(exc: OSError) -> int: code = exc.winerror # type: ignore[attr-defined] if isinstance(code, int) and code: @@ -358,13 +423,19 @@ def _fail(api: str, code: int | None = None) -> NoReturn: class Job: - """A job handle owned by the parent process.""" + """Retain the job and identity-pinned members until kernel exit.""" def __init__(self, handle: ctypes.c_void_p, name: str) -> None: self._handle: ctypes.c_void_p | None = handle self.name = name + self._process_handles: dict[int, ctypes.c_void_p] = {} def close(self) -> None: + while self._process_handles: + pid = next(iter(self._process_handles)) + if not _kernel32().CloseHandle(self._process_handles[pid]): + raise _win_os_error(_last_error()) + del self._process_handles[pid] handle = self._handle if handle is not None and not _kernel32().CloseHandle(handle): raise _win_os_error(_last_error()) @@ -394,12 +465,98 @@ def active_processes(self) -> int | None: return None return int(info.ActiveProcesses) + def _seal_admission(self) -> None: + limits = ExtendedLimitInformation() + handle = self._open_handle() + if not _kernel32().QueryInformationJobObject( + handle, + _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, + ctypes.byref(limits), + ctypes.sizeof(limits), + None, + ): + raise _win_os_error(_last_error()) + # No member may fork between enumeration and termination. + limits.BasicLimitInformation.LimitFlags |= _JOB_OBJECT_LIMIT_ACTIVE_PROCESS + limits.BasicLimitInformation.ActiveProcessLimit = 0 + if not _kernel32().SetInformationJobObject( + handle, + _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, + ctypes.byref(limits), + ctypes.sizeof(limits), + ): + raise _win_os_error(_last_error()) + + def _process_ids(self) -> list[int]: + capacity = 1 + # Sealed admission permits one size query and one complete snapshot. + for _attempt in range(_JOB_LIST_QUERIES): + buffer = ctypes.create_string_buffer( + ProcessIdList.ProcessIdList.offset + + capacity * ctypes.sizeof(ctypes.c_size_t) + ) + info = ctypes.cast(buffer, ctypes.POINTER(ProcessIdList)).contents + succeeded = _kernel32().QueryInformationJobObject( + self._open_handle(), + _JOB_OBJECT_BASIC_PROCESS_ID_LIST, + buffer, + ctypes.sizeof(buffer), + None, + ) + if not succeeded and (code := _last_error()) != _ERROR_MORE_DATA: + raise _win_os_error(code) + assigned = int(info.NumberOfAssignedProcesses) + count = int(info.NumberOfProcessIdsInList) + if succeeded and count == assigned and count <= capacity: + pids = (ctypes.c_size_t * count).from_buffer( + buffer, ProcessIdList.ProcessIdList.offset + ) + return list(pids) + capacity = max(capacity * 2, assigned) + raise _win_os_error(_ERROR_MORE_DATA) + + def _capture_process_handles(self) -> None: + for pid in self._process_ids(): + if pid in self._process_handles: + continue + handle = _kernel32().OpenProcess( + _SYNCHRONIZE | _PROCESS_QUERY_LIMITED_INFORMATION, False, pid + ) + if not handle: + code = _last_error() + if code == _ERROR_INVALID_PARAMETER: + continue + raise _win_os_error(code) + self._process_handles[pid] = ctypes.c_void_p(handle) + for pid, handle in tuple(self._process_handles.items()): + member = wintypes.BOOL() + if not _kernel32().IsProcessInJob( + handle, self._open_handle(), ctypes.byref(member) + ): + raise _win_os_error(_last_error()) + if not member.value: + # Opening a PID can race exit and reuse; never adopt that identity. + if not _kernel32().CloseHandle(handle): + raise _win_os_error(_last_error()) + del self._process_handles[pid] + + def _pending_process_exits(self) -> int: + pending = 0 + for handle in self._process_handles.values(): + state = _kernel32().WaitForSingleObject(handle, 0) + if state == _WAIT_FAILED: + raise _win_os_error(_last_error()) + pending += state != _WAIT_OBJECT_0 + return pending + def terminate(self) -> None: """Mark every member for termination. Members blocked in a synchronous pipe read stay inside that call until the other end closes. Wait only after that end is closed. """ + self._seal_admission() + self._capture_process_handles() if not _kernel32().TerminateJobObject(self._open_handle(), 1): raise _win_os_error(_last_error()) @@ -409,11 +566,16 @@ async def wait_until_idle(self, timeout: float) -> None: deadline = loop.time() + timeout while True: active = self.active_processes() - if active == 0: + pending = self._pending_process_exits() + if active == 0 and not pending: return remaining = deadline - loop.time() if active is None or remaining <= 0: - logger.warning("Codex job still has %s active processes", active) + logger.warning( + "Codex job has %s active processes and %s pending kernel exits", + active, + pending, + ) raise JobNotDrained(active) await asyncio.sleep(min(_JOB_POLL_S, remaining)) @@ -428,6 +590,7 @@ def create_job() -> Job: _kernel32().CloseHandle(handle) _fail("CreateJobObject", _ERROR_ALREADY_EXISTS) limits = ExtendedLimitInformation() + # Breakaway flags stay unset so descendants cannot leave this job. limits.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE if not _kernel32().SetInformationJobObject( handle, @@ -529,6 +692,44 @@ async def await_returncode(process: asyncio.subprocess.Process, timeout: float) return True +def session_pid_reused(process: asyncio.subprocess.Process) -> bool: + """True when ``process.pid`` now names some other live process. + + After the leader is reaped, its session id stays reserved only while a + member of that session still exists. Once the id is free it can be given + to a new session, and ``killpg`` would signal that session. + """ + if process.returncode is None: + return False + try: + os.kill(process.pid, 0) + except ProcessLookupError: + return False + except OSError: + logger.warning( + "Process %s pid is not signalable; leaving its session", + process.pid, + ) + return True + logger.warning( + "Process %s pid was reused; not signaling that session", + process.pid, + ) + return True + + +async def kill_process_tree(process: asyncio.subprocess.Process) -> None: + """Force-stop the owned Unix process group, including after root exit. + + On Unix the launched child leads a session, so ``killpg`` still reaches + that group after the leader has exited. + """ + if os.name == "nt" or session_pid_reused(process): + return + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + + async def close_stdin(process: asyncio.subprocess.Process) -> None: """Drop the write end so a synchronous read in the child can finish. @@ -550,10 +751,11 @@ async def _launch_helper( cwd: str | None, env: Mapping[str, str] | None, limit: int, + create_process: CreateProcess, ) -> asyncio.subprocess.Process: startup = subprocess.STARTUPINFO() # type: ignore[attr-defined] startup.lpAttributeList = {"handle_list": [status_handle]} - return await asyncio.create_subprocess_exec( + return await create_process( *command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, @@ -568,7 +770,8 @@ async def _launch_helper( class WindowsProcess: """One client-owned helper and job, including incomplete startup.""" - def __init__(self) -> None: + def __init__(self, *, create_process: CreateProcess | None = None) -> None: + self.create_process = create_process or asyncio.create_subprocess_exec self.job: Job | None = create_job() self.process: asyncio.subprocess.Process | None = None self.closing = False @@ -576,6 +779,22 @@ def __init__(self) -> None: self._status: asyncio.Task[str] | None = None self._read_fd: int | None = None + async def start_command( + self, + command: Sequence[str], + *, + cwd: str | None, + env: Mapping[str, str] | None, + limit: int, + ) -> asyncio.subprocess.Process: + """Start a command with job membership established before it can run.""" + return await self.start( + functools.partial(command_in_job, command), + cwd=cwd, + env=env, + limit=limit, + ) + async def start( self, command_for: Callable[[str, int], Sequence[str]], @@ -622,6 +841,7 @@ async def _create( cwd=cwd, env=env, limit=limit, + create_process=self.create_process, ) finally: # Only the helper may keep the writer; its exit must produce EOF. diff --git a/src/band/integrations/omp.py b/src/band/integrations/omp.py index a4370d697..1e0e189da 100644 --- a/src/band/integrations/omp.py +++ b/src/band/integrations/omp.py @@ -24,6 +24,12 @@ # The acp subcommand's flag for the room workspace, passed as ``--cwd=``. OMP_CWD_FLAG = "--cwd" + +def omp_cwd_argument(workspace: str) -> str: + """Format OMP's workspace flag without changing launch cwd.""" + return f"{OMP_CWD_FLAG}={workspace}" + + DEFAULT_OMP_ACP_COMMAND: tuple[str, ...] = ( "omp", OMP_ACP_SUBCOMMAND, @@ -167,7 +173,7 @@ def omp_command_in_workspace(command: Sequence[str], workspace: str) -> list[str f"got {list(command)!r}" ) split = list(command).index(OMP_ACP_SUBCOMMAND) + 1 - return [*command[:split], f"{OMP_CWD_FLAG}={workspace}", *command[split:]] + return [*command[:split], omp_cwd_argument(workspace), *command[split:]] def omp_elicitation_call_id(session_id: str) -> str: @@ -210,12 +216,12 @@ def is_omp_approve_deny_form(requested_schema: object) -> bool: return approve_deny_form_field(requested_schema) is not None -def _omp_mcp_remainder(name_or_path: str) -> str | None: - """Strip ``xd://`` and/or ``mcp__``; ``None`` when not an OMP MCP spelling.""" +def _omp_mcp_remainder(name_or_path: str, namespace: str) -> str | None: + """Strip ``xd://``, ``mcp__`` and ``namespace``; ``None`` when not an OMP MCP spelling.""" text = name_or_path.removeprefix(XD_URL_PREFIX) if not text.startswith(XD_MCP_PREFIX): return None - remainder = text.removeprefix(XD_MCP_PREFIX) + remainder = text.removeprefix(XD_MCP_PREFIX).removeprefix(namespace) return remainder or None @@ -243,10 +249,14 @@ def _omp_mcp_wire_candidates(remainder: str) -> list[str]: def normalize_omp_mcp_tool_name( - name_or_path: str, own_names: Collection[str] + name_or_path: str, own_names: Collection[str], *, namespace: str = "" ) -> str | None: - """Map an OMP ``mcp__…`` title or ``xd://mcp__…`` path to a Band tool name.""" - remainder = _omp_mcp_remainder(name_or_path) + """Map an OMP ``mcp__…`` title or ``xd://mcp__…`` path to a Band tool name. + + ``namespace`` is a server-name prefix OMP mints into every tool of a + plugin-projected server; it is removed before the remainder is matched. + """ + remainder = _omp_mcp_remainder(name_or_path, namespace) if remainder is None: return None for candidate in _omp_mcp_wire_candidates(remainder): @@ -262,6 +272,7 @@ def normalize_omp_mcp_device_call( own_names: Collection[str], *, kind: str | None = None, + namespace: str = "", ) -> tuple[str, dict[str, object]]: """Map OMP MCP device writes / ``mcp__`` titles to canonical Band tool names. @@ -281,7 +292,7 @@ def normalize_omp_mcp_device_call( break if path is not None: - canonical = normalize_omp_mcp_tool_name(path, own_names) + canonical = normalize_omp_mcp_tool_name(path, own_names, namespace=namespace) content = args.get("content") if canonical is None or kind == _ACP_KIND_READ: return name, args @@ -297,7 +308,7 @@ def normalize_omp_mcp_device_call( return canonical, args return canonical, parsed - canonical = normalize_omp_mcp_tool_name(name, own_names) + canonical = normalize_omp_mcp_tool_name(name, own_names, namespace=namespace) if canonical is None: return name, args return canonical, args diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py new file mode 100644 index 000000000..eb0ed34ae --- /dev/null +++ b/src/band/integrations/ompprocess.py @@ -0,0 +1,343 @@ +"""Native OMP (``omp --mode rpc``) launch settings and the SDK-owned MCP plugin.""" + +from __future__ import annotations + +import json +import os +import tempfile +from collections.abc import Collection, Iterator, Mapping, Sequence +from contextlib import contextmanager +from enum import StrEnum +from pathlib import Path +from typing import TYPE_CHECKING, Literal, Self +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from band.integrations.omp import ( + normalize_omp_mcp_device_call, + normalize_omp_mcp_tool_name, + omp_cwd_argument, + omp_provider_env, +) +from band.integrations.ompwire import ( + OmpCommandName, + OmpEvent, + OmpEventType, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +if TYPE_CHECKING: + from band.integrations.omprpc import OmpRpcClient + +OMP_MODE_FLAG = "--mode" +OMP_RPC_MODE = "rpc" +OMP_RPC_UI_MODE = "rpc-ui" +DEFAULT_OMP_COMMAND: tuple[str, ...] = ("omp", OMP_MODE_FLAG, OMP_RPC_MODE) +RELATIVE_MCP_COMMAND_PREFIXES = ("./", "../", ".\\", "..\\") + +# The native RPC ask dialog and MCP status events need a newer OMP than the +# ACP bridge's pin. +OMP_NATIVE_PINNED_PACKAGE = "@oh-my-pi/pi-coding-agent@18.8.9" + +# Plugin names prefix the MCP server names OMP mints, so the plugin root's +# basename must stay stable across starts; its sanitized form is the +# namespace every projected tool name carries. +OMP_PLUGIN_ROOT = "band-sdk" +OMP_PLUGIN_TOOL_NAMESPACE = f"{OMP_PLUGIN_ROOT.replace('-', '_')}_" +OMP_MCP_STARTUP_JS = "startup.js" +OMP_MCP_STARTUP_CONFIG = "startup.json" +OMP_MCP_RECEIPT = "receipt.json" +OMP_MCP_SERVERS_CONFIG = ".mcp.json" + + +class McpConnectionStatus(StrEnum): + CONNECTED = "connected" + FAILED = "failed" + + +def omp_runtime_env( + *, + env: Mapping[str, str] | None, + model: str | None, + api_key: str | None, +) -> dict[str, str] | None: + """Build the provider environment, preserving explicitly supplied overrides.""" + credentials = ( + omp_provider_env(model=model, api_key=api_key) + if model and api_key is not None + else {} + ) + return {**credentials, **(env or {})} or None + + +def omp_rpc_command_in_workspace(command: Sequence[str], workspace: str) -> list[str]: + """Select OMP's room workspace without changing the subprocess launch cwd.""" + return [*command, omp_cwd_argument(workspace)] + + +class OmpMcpTransport(StrEnum): + """The transports OMP accepts for external MCP servers.""" + + STDIO = "stdio" + HTTP = "http" + SSE = "sse" + + +class OmpMcpServer(BaseModel): + """An external MCP server projected into the SDK-owned OMP plugin.""" + + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + type: OmpMcpTransport | None = None + command: str | None = None + args: tuple[str, ...] = () + env: dict[str, str] = Field(default_factory=dict, repr=False) + cwd: str | None = None + url: str | None = None + headers: dict[str, str] = Field(default_factory=dict, repr=False) + + @model_validator(mode="after") + def _require_one_transport(self) -> Self: + if bool(self.command) == bool(self.url): + raise ValueError("MCP server needs exactly one of command or url") + if self.command and self.type not in (None, OmpMcpTransport.STDIO): + raise ValueError(f"a command requires {OmpMcpTransport.STDIO} transport") + if self.url and self.type is OmpMcpTransport.STDIO: + raise ValueError( + f"a URL requires {OmpMcpTransport.HTTP} or {OmpMcpTransport.SSE} transport" + ) + return self + + +_MCP_STARTUP_EXTENSION = """ +import { writeFileSync } from "node:fs"; + +export default function (pi) { + const expected = new Set(config.servers); + const statuses = new Map(); + let finish; + const ready = new Promise(resolve => { finish = resolve; }); + const unsubscribe = pi.events.on("mcp:connection-status", event => { + if (!expected.has(event.serverName) || statuses.has(event.serverName)) return; + if (event.type !== CONNECTED && event.type !== FAILED) return; + statuses.set(event.serverName, event); + if (event.type === FAILED || statuses.size === expected.size) { + unsubscribe(); + finish(); + } + }); + if (expected.size === 0) { + unsubscribe(); + finish(); + } + pi.registerCommand(config.command, { + description: "Verify the host-configured MCP connections", + handler: async (_args, ctx) => { + await ready; + const trimmed = (config.systemPrompt || "").trim(); + const systemPromptReady = !trimmed + || ctx.getSystemPrompt().some(block => block.includes(trimmed)); + writeFileSync(config.receipt, JSON.stringify({ systemPromptReady, servers: [...statuses.values()] }), { mode: 0o600 }); + }, + }); +} +""" + + +class OmpStartupError(RuntimeError): + """The required native connection context or MCP tools are not ready.""" + + +class McpCommand(BaseModel): + name: str + + +class McpCommands(BaseModel): + commands: list[McpCommand] + + +class ConnectedMcpServer(BaseModel): + type: Literal[McpConnectionStatus.CONNECTED] + serverName: str + + +class FailedMcpServer(BaseModel): + type: Literal[McpConnectionStatus.FAILED] + serverName: str + error: str + + +class McpStartupReceipt(BaseModel): + systemPromptReady: bool = Field(strict=True) + servers: list[ConnectedMcpServer | FailedMcpServer] + + +def _plugin_server_name(name: str) -> str: + return f"{OMP_PLUGIN_ROOT}:{name}" + + +def omp_mcp_ready_command(plugin: str) -> str: + """The slash command the plugin's startup extension registers for its receipt.""" + return f"{OMP_PLUGIN_ROOT}-mcp-{Path(plugin).parent.name}" + + +def omp_mcp_arguments(plugin: str) -> tuple[str, ...]: + """Enable native MCP status events without changing the user's persisted settings.""" + return ( + "--plugin-dir", + plugin, + OMP_MODE_FLAG, + OMP_RPC_UI_MODE, + "--extension", + str(Path(plugin) / OMP_MCP_STARTUP_JS), + "--config", + str(Path(plugin) / OMP_MCP_STARTUP_CONFIG), + ) + + +async def omp_mcp_ready( + runtime: OmpRpcClient, + plugin: str, + servers: Mapping[str, OmpMcpServer], +) -> None: + """Require the connection extension and configured MCP servers before inference.""" + command = omp_mcp_ready_command(plugin) + commands = McpCommands.model_validate( + await runtime.request(OmpCommandName.GET_AVAILABLE_COMMANDS) + ) + if not any(item.name == command for item in commands.commands): + raise OmpStartupError("OMP did not load the required SDK startup extension") + + async def on_event(event: OmpEvent) -> None: + if event.type == OmpEventType.AGENT_START: + raise OmpStartupError( + "OMP unexpectedly invoked an agent during SDK startup" + ) + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + raise OmpStartupError("OMP unexpectedly invoked a host tool during SDK startup") + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + raise OmpStartupError("OMP unexpectedly requested UI during SDK startup") + + await runtime.prompt( + f"/{command}", + on_event=on_event, + on_tool=on_tool, + on_ui=on_ui, + ) + try: + receipt = McpStartupReceipt.model_validate_json( + (Path(plugin) / OMP_MCP_RECEIPT).read_text(encoding="utf-8") + ) + except (OSError, ValueError) as error: + raise OmpStartupError("OMP returned no valid SDK startup receipt") from error + if not receipt.systemPromptReady: + raise OmpStartupError( + "OMP did not load the required Band connection instructions; " + "rules may be disabled or omitted by the system prompt template" + ) + expected = {_plugin_server_name(name) for name in servers} + for status in receipt.servers: + if status.serverName not in expected: + raise OmpStartupError("OMP returned an unexpected MCP startup receipt") + if isinstance(status, FailedMcpServer): + raise OmpStartupError( + f"OMP MCP server {status.serverName} failed: {status.error}" + ) + if {status.serverName for status in receipt.servers} != expected or len( + receipt.servers + ) != len(expected): + raise OmpStartupError("OMP did not acknowledge every configured MCP server") + + +def _projected_server(server: OmpMcpServer, workspace: str) -> dict[str, object]: + payload = server.model_dump(mode="json", exclude_none=True, exclude_defaults=True) + if server.command: + # Plugin discovery otherwise rebases launch paths into the temporary overlay. + cwd = Path(os.path.abspath(Path(workspace) / (server.cwd or "."))) + payload["cwd"] = str(cwd) + if server.command.startswith(RELATIVE_MCP_COMMAND_PREFIXES): + payload["command"] = os.path.abspath(cwd / server.command) + return payload + + +def _write_private(path: Path, content: str) -> None: + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(content) + + +@contextmanager +def omp_mcp_plugin( + servers: Mapping[str, OmpMcpServer], *, workspace: str, system_prompt: str = "" +) -> Iterator[str | None]: + """Project connection instructions and servers without changing native user configs.""" + trimmed_prompt = system_prompt.strip() + if not servers and not trimmed_prompt: + yield None + return + projected = { + name: _projected_server(server, workspace) for name, server in servers.items() + } + with tempfile.TemporaryDirectory(prefix="band-omp-") as temporary: + root = Path(temporary) / OMP_PLUGIN_ROOT + root.mkdir(mode=0o700) + if trimmed_prompt: + rules = root / "rules" + rules.mkdir(mode=0o700) + _write_private( + rules / f"band-connection-{uuid4().hex}.md", + f"---\nalwaysApply: true\n---\n{trimmed_prompt}", + ) + startup = { + "servers": [_plugin_server_name(name) for name in servers], + "command": omp_mcp_ready_command(str(root)), + "receipt": str(root / OMP_MCP_RECEIPT), + "systemPrompt": trimmed_prompt, + } + _write_private( + root / OMP_MCP_SERVERS_CONFIG, json.dumps({"mcpServers": projected}) + ) + # Quiet startup suppresses the native status bus, not just UI rendering. + _write_private( + root / OMP_MCP_STARTUP_CONFIG, json.dumps({"startup": {"quiet": False}}) + ) + _write_private( + root / OMP_MCP_STARTUP_JS, + "\n".join( + ( + f"const config = {json.dumps(startup)};", + f"const CONNECTED = {json.dumps(McpConnectionStatus.CONNECTED)};", + f"const FAILED = {json.dumps(McpConnectionStatus.FAILED)};", + _MCP_STARTUP_EXTENSION, + ) + ), + ) + yield str(root) + + +def normalize_native_mcp_tool_name( + name_or_path: str, own_names: Collection[str] +) -> str | None: + """Map a tool name OMP minted under the SDK plugin namespace to a Band tool.""" + return normalize_omp_mcp_tool_name( + name_or_path, own_names, namespace=OMP_PLUGIN_TOOL_NAMESPACE + ) + + +def normalize_native_mcp_device_call( + name: str, + arguments: Mapping[str, object], + own_names: Collection[str], + *, + kind: str | None = None, +) -> tuple[str, dict[str, object]]: + """Map a device call under the SDK plugin namespace to a canonical Band tool.""" + return normalize_omp_mcp_device_call( + name, arguments, own_names, kind=kind, namespace=OMP_PLUGIN_TOOL_NAMESPACE + ) diff --git a/src/band/integrations/omprpc/__init__.py b/src/band/integrations/omprpc/__init__.py new file mode 100644 index 000000000..5df95e812 --- /dev/null +++ b/src/band/integrations/omprpc/__init__.py @@ -0,0 +1,19 @@ +"""Persistent subprocess transport for OMP's native JSONL RPC protocol.""" + +from __future__ import annotations + +from band.integrations.omprpc.client import OmpRpcClient +from band.integrations.omprpc.errors import ( + OmpProcessError, + OmpPromptError, + OmpProtocolError, + OmpRpcError, +) + +__all__ = [ + "OmpProcessError", + "OmpPromptError", + "OmpProtocolError", + "OmpRpcClient", + "OmpRpcError", +] diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py new file mode 100644 index 000000000..dec072fab --- /dev/null +++ b/src/band/integrations/omprpc/client.py @@ -0,0 +1,869 @@ +"""Persistent subprocess transport for OMP's native JSONL RPC protocol.""" + +from __future__ import annotations + +import asyncio +import logging +import os +from binascii import Error as Base64Error +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, TypeVar, cast + +from pydantic import JsonValue, ValidationError + +from band.integrations.codex.winspawn import ( + CreateProcess, + SpawnFailure, + WindowsProcess, + await_returncode, + close_stdin, + kill_process_tree, +) +from band.integrations.omprpc.errors import ( + CALLBACKS_UNFINISHED, + OmpProcessError, + OmpPromptError, + OmpProtocolError, + OmpRpcError, + command_error, +) +from band.integrations.omprpc.frames import ( + ChunkAssembler, + parse_payload_line, + payload_json, + ui_response_frame, +) +from band.integrations.omprpc.prompt import OmpCallbackKind, PromptRun +from band.integrations.omprpc.scope import OmpProcessScope +from band.integrations.ompwire import ( + OMP_FRAME_BYTES, + OMP_NATIVE_PROTOCOL_VERSION, + OMP_REASSEMBLED_BYTES, + OmpCommand, + OmpCommandName, + OmpEvent, + OmpEventType, + OmpFrame, + OmpFrameError, + OmpFrameType, + OmpHostToolCall, + OmpHostToolCancel, + OmpHostToolResult, + OmpNegotiation, + OmpPromptAcceptance, + OmpPromptResult, + OmpPromptStatus, + OmpReady, + OmpResponse, + OmpState, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, + OmpWireModel, +) + +TModel = TypeVar("TModel", bound=OmpWireModel) +CloseProcess = Callable[[], Coroutine[Any, Any, None]] +JoinClose = Callable[[Awaitable[None]], Awaitable[None]] + +logger = logging.getLogger(__name__) + + +def _noop() -> None: + return None + + +_DRAIN_BYTES = 64 * 1024 +_ABORT_TIMEOUT_S = 1.0 +PROCESS_EXIT_TIMEOUT_S = 1.0 +CALLBACK_EXIT_TIMEOUT_S = 1.0 +_EVENT_TYPES = frozenset(OmpEventType) +_UI_NOTIFICATIONS = frozenset( + {"notify", "setStatus", "setWidget", "setTitle", "set_editor_text", "open_url"} +) + + +@dataclass(frozen=True) +class ReceivedResponse: + frame: OmpResponse + sequence: int + + +@dataclass(frozen=True) +class PendingRequest: + command: str + future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] + + +class OmpRpcClient: + """One OMP process, one persistent reader, and at most one active prompt.""" + + def __init__( + self, + command: Sequence[str], + *, + env: Mapping[str, str] | None = None, + create_process: CreateProcess | None = None, + ) -> None: + if not command: + raise ValueError("OMP command must not be empty") + self.command = tuple(command) + self.env = None if env is None else {**os.environ, **env} + self.create_process = create_process + self.close_process: CloseProcess = self._close_process + self.join_close: JoinClose = asyncio.shield + self.after_stdin_closed: Callable[[], None] = _noop + self.callback_exit_timeout_s = CALLBACK_EXIT_TIMEOUT_S + self.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S + self._proc: asyncio.subprocess.Process | None = None + self._windows_process: WindowsProcess | None = None + self._process_scope: OmpProcessScope | None = None + self._spawn_task: asyncio.Task[asyncio.subprocess.Process] | None = None + self._reader_task: asyncio.Task[None] | None = None + self._stderr_task: asyncio.Task[None] | None = None + self._close_task: asyncio.Task[None] | None = None + self._ready: asyncio.Future[OmpReady | OmpProcessError] | None = None + self._failure: OmpProcessError | None = None + self._closing = False + self._request_id = 0 + self._sequence = 0 + self._frames = ChunkAssembler() + self._max_frame_bytes = OMP_FRAME_BYTES + self._pending: dict[str, PendingRequest] = {} + self._turn: PromptRun | None = None + self._start_lock = asyncio.Lock() + self._write_lock = asyncio.Lock() + self._prompt_lock = asyncio.Lock() + self._write_tasks: set[asyncio.Task[None]] = set() + + async def start(self) -> None: + """Start the supplied command and require the native RPC protocol.""" + async with self._start_lock: + if self._closing: + await self.close() + self._proc = None + self._spawn_task = None + self._close_task = None + if self._proc is not None: + self._ensure_running() + return + if self._failure is not None: + logger.error( + "Discarding previous OMP RPC failure on restart: %s", self._failure + ) + self._failure = None + self._frames.reset() + self._frames.max_reassembled_bytes = OMP_REASSEMBLED_BYTES + self._max_frame_bytes = OMP_FRAME_BYTES + self._closing = False + self._process_scope = OmpProcessScope() + self._ready = asyncio.get_running_loop().create_future() + self._spawn_task = asyncio.create_task(self._spawn_process()) + try: + process = await asyncio.shield(self._spawn_task) + if self._closing: + raise OmpProcessError("OMP RPC client closed during startup") + self._proc = process + self._reader_task = asyncio.create_task(self._read_stdout(self._proc)) + self._stderr_task = asyncio.create_task(self._read_stderr(self._proc)) + ready = await asyncio.shield(self._ready) + if isinstance(ready, OmpProcessError): + raise ready + if OMP_NATIVE_PROTOCOL_VERSION not in ready.supportedProtocolVersions: + error = OmpProtocolError( + f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} is required" + ) + self._record_failure(error) + raise error + self._apply_ready_limits(ready) + negotiated = await self._request_frame( + OmpCommandName.NEGOTIATE_PROTOCOL, + {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, + ) + self._parse_response( + negotiated.frame, + OmpNegotiation, + f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} negotiation failed", + ) + except (OSError, SpawnFailure) as cause: + await self.close() + raise OmpProcessError("Unable to start OMP RPC process") from cause + except (Exception, asyncio.CancelledError): + await self.close() + raise + + async def spawn_owned_process( + self, *command: str, **kwargs: Any + ) -> asyncio.subprocess.Process: + """Start and retain the command within this client's process containment. + + ``create_process`` is the low-level launcher, including for the Windows + helper. Wrappers override this method and delegate to its bound original + or ``super()`` so creation owns exactly one helper and Job. + """ + if os.name == "nt": + self._windows_process = WindowsProcess(create_process=self.create_process) + process = await self._windows_process.start_command( + command, + cwd=cast(str | None, kwargs.get("cwd")), + env=cast(Mapping[str, str] | None, kwargs.get("env")), + limit=cast(int, kwargs.get("limit", OMP_FRAME_BYTES)), + ) + else: + create = self.create_process or asyncio.create_subprocess_exec + process = await create(*command, **kwargs) + # A wrapper may hold or fail after creation; cleanup still owns the child. + self._proc = process + return process + + async def _spawn_process(self) -> asyncio.subprocess.Process: + scope = self._process_scope + assert scope is not None + env = scope.environment(self.env) + return await self.spawn_owned_process( + *self.command, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=None, + env=env, + start_new_session=os.name != "nt", + limit=OMP_FRAME_BYTES, + ) + + async def request( + self, command: str, **parameters: JsonValue + ) -> dict[str, JsonValue]: + """Correlate a native command response without consuming prompt events.""" + self._ensure_running() + response = await self._request_frame(command, parameters) + return self._response_data(response.frame) + + async def prompt( + self, + message: str, + *, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + """Deliver a turn through its correlated result and actual native settlement.""" + async with self._prompt_lock: + self._ensure_running() + await self._execute_prompt(message, on_event, on_tool, on_ui) + + async def close(self) -> None: + """Finish shared bounded cleanup before propagating caller cancellation.""" + if self._proc is None and self._spawn_task is None: + return + self._begin_close() + closing = self._close_task + assert closing is not None + interrupted = False + while True: + try: + await self.join_close(closing) + break + except asyncio.CancelledError: + if closing.cancelled(): + raise + interrupted = True + if interrupted: + raise asyncio.CancelledError + + def _ensure_running(self) -> None: + if self._failure is not None: + raise self._failure + if self._closing or self._proc is None: + raise OmpProcessError("OMP RPC process is not running") + if self._proc.returncode is not None: + error = OmpProcessError("OMP RPC process exited") + self._record_failure(error) + raise error + + def _next_id(self) -> str: + self._request_id += 1 + return f"band-{self._request_id}" + + async def _request_frame( + self, + command: str, + parameters: Mapping[str, JsonValue], + *, + request_id: str | None = None, + allow_closing: bool = False, + ) -> ReceivedResponse: + if not allow_closing: + self._ensure_running() + if "id" in parameters or "type" in parameters: + raise ValueError("OMP command parameters must not override id or type") + request_id = request_id or self._next_id() + frame = OmpCommand.model_validate( + {"type": str(command), "id": request_id, **parameters} + ) + future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] = ( + asyncio.get_running_loop().create_future() + ) + self._pending[request_id] = PendingRequest(command, future) + try: + await self._write_frame(frame.root) + response = await future + if isinstance(response, (OmpRpcError, OmpProcessError)): + raise response + return response + finally: + self._pending.pop(request_id, None) + + def _response_data(self, response: OmpResponse) -> dict[str, JsonValue]: + if response.data is None: + return {} + if not isinstance(response.data, dict): + error = OmpProtocolError("OMP command response data must be an object") + self._record_failure(error) + raise error + return response.data + + def _parse_response( + self, response: OmpResponse, model: type[TModel], error: str + ) -> TModel: + try: + return model.model_validate(self._response_data(response)) + except ValidationError as cause: + failure = OmpProtocolError(error) + failure.__cause__ = cause + self._record_failure(failure) + raise failure from cause + + async def _execute_prompt( + self, + message: str, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + turn = PromptRun(self._next_id(), on_event, on_tool, on_ui) + self._turn = turn + event_worker = asyncio.create_task(self._deliver_events(turn)) + turn.event_worker = event_worker + work = asyncio.create_task(self._complete_prompt(turn, message)) + try: + done, _ = await asyncio.wait( + (work, turn.failure), return_when=asyncio.FIRST_COMPLETED + ) + if turn.failure in done: + raise turn.failure.result() + await work + except OmpRpcError: + if not turn.rejected: + await self.close() + raise + except OmpPromptError: + raise + except BaseException: + await self.close() + raise + finally: + self._turn = None + await self._cancel_tasks( + [work, event_worker, *turn.callback_tasks], + timeout=self.callback_exit_timeout_s, + label="callback", + ) + + async def _complete_prompt(self, turn: PromptRun, message: str) -> None: + try: + accepted = await self._request_frame( + OmpCommandName.PROMPT, {"message": message}, request_id=turn.id + ) + except OmpRpcError: + turn.rejected = True + raise + acceptance = self._parse_response( + accepted.frame, + OmpPromptAcceptance, + "Invalid OMP prompt acceptance protocol data", + ) + outcome: OmpPromptResult | None = None + # A synchronous local command explicitly owes no prompt_result. + if acceptance.agentInvoked is not False: + outcome = await turn.result + await self._wait_for_settlement(turn) + await turn.events.join() + if turn.callback_tasks: + try: + async with asyncio.timeout(self.callback_exit_timeout_s): + while any(not task.done() for task in turn.callback_tasks): + await asyncio.wait(turn.callback_tasks) + except TimeoutError: + error = OmpProcessError(CALLBACKS_UNFINISHED) + self._record_failure(error) + raise error from None + if outcome is not None and outcome.status != OmpPromptStatus.COMPLETED: + raise OmpPromptError( + outcome.status, + outcome.error.message + if outcome.error is not None + else f"OMP prompt {outcome.status}", + ) + assistant = turn.last_assistant + if assistant is not None and assistant.stopReason in { + OmpPromptStatus.ERROR, + OmpPromptStatus.ABORTED, + }: + raise OmpPromptError( + assistant.stopReason, + assistant.errorMessage or f"OMP continuation {assistant.stopReason}", + ) + + async def _wait_for_settlement(self, turn: PromptRun) -> None: + snapshot = await self._request_frame(OmpCommandName.GET_STATE, {}) + state = self._parse_response( + snapshot.frame, OmpState, "Invalid OMP settlement protocol data" + ) + if state.isSettled and snapshot.sequence >= turn.last_activity_sequence: + return + # An old, uncorrelated session_settled cannot overrule a newer busy snapshot. + while ( + turn.settled_sequence <= snapshot.sequence + or turn.settled_sequence <= turn.last_activity_sequence + ): + turn.changed.clear() + await turn.changed.wait() + + async def _deliver_events(self, turn: PromptRun) -> None: + try: + while True: + event = await turn.events.get() + try: + await turn.on_event(event) + finally: + turn.events.task_done() + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 -- arbitrary callback failures belong to the awaiting prompt + turn.fail(error) + + async def _write_line(self, line: str) -> None: + async with self._write_lock: + process = self._proc + if process is None or process.stdin is None: + raise OmpProcessError("OMP RPC stdin is unavailable") + try: + process.stdin.write((line + "\n").encode("utf-8")) + await process.stdin.drain() + except (BrokenPipeError, ConnectionError, OSError) as cause: + error = OmpProcessError("OMP RPC stdin closed") + error.__cause__ = cause + self._record_failure(error) + raise error from cause + + async def _write_frame(self, frame: dict[str, JsonValue]) -> None: + await self._write_line(payload_json(frame)) + + async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: + assert process.stdout is not None + try: + while True: + line = await process.stdout.readline() + if not line: + if self._closing: + return + if self._frames.pending: + raise OmpProtocolError( + "OMP RPC protocol chunks interrupted by stdout EOF" + ) + raise OmpProcessError("OMP RPC stdout EOF") + if len(line) > self._max_frame_bytes or not line.endswith(b"\n"): + raise OmpProtocolError("Invalid OMP RPC protocol frame length") + payload = self._frames.decode(parse_payload_line(line)) + if payload is None: + continue + self._sequence += 1 + await self._dispatch_payload(payload, self._sequence) + except asyncio.CancelledError: + raise + except OmpProcessError as error: + self._record_failure(error) + except (ValidationError, ValueError, Base64Error) as cause: + error = OmpProtocolError("Invalid OMP RPC protocol frame") + error.__cause__ = cause + self._record_failure(error) + except Exception as cause: # noqa: BLE001 -- reader failure must unblock every pending request + error = OmpProcessError("OMP RPC stdout reader failed") + error.__cause__ = cause + self._record_failure(error) + + async def _dispatch_payload( + self, payload: dict[str, JsonValue], sequence: int + ) -> None: + frame = OmpFrame.model_validate(payload) + turn = self._turn + match frame.type: + case OmpFrameType.READY: + ready = OmpReady.model_validate(payload) + self._apply_ready_limits(ready) + if self._ready is not None and not self._ready.done(): + self._ready.set_result(ready) + else: + logger.debug("OMP RPC ignored a duplicate ready frame") + case OmpFrameType.FRAME_ERROR: + overflow = OmpFrameError.model_validate(payload) + raise OmpProtocolError( + f"OMP RPC protocol frame overflow: {overflow.error}" + ) + case OmpFrameType.RESPONSE: + self._dispatch_response(OmpResponse.model_validate(payload), sequence) + case OmpFrameType.PROMPT_RESULT: + result = OmpPromptResult.model_validate(payload) + if turn is not None and result.id == turn.id and not turn.result.done(): + turn.result.set_result(result) + else: + logger.debug("OMP RPC ignored an unmatched prompt_result") + case OmpFrameType.SESSION_SETTLED: + if turn is not None and ( + turn.last_activity_sequence > 0 or turn.result.done() + ): + turn.settled_sequence = sequence + turn.changed.set() + case OmpFrameType.HOST_TOOL_CALL: + call = OmpHostToolCall.model_validate(payload) + if turn is None: + raise OmpProtocolError( + "OMP host tool arrived without an active prompt" + ) + self._start_callback( + turn, OmpCallbackKind.TOOL, call.id, self._handle_tool(turn, call) + ) + case OmpFrameType.HOST_TOOL_CANCEL: + cancellation = OmpHostToolCancel.model_validate(payload) + if turn is not None: + self._cancel_callback( + turn, OmpCallbackKind.TOOL, cancellation.targetId + ) + else: + logger.debug( + "OMP RPC ignored a host-tool cancellation without an active prompt" + ) + case OmpFrameType.UI_REQUEST: + request = OmpUIRequest.model_validate(payload) + if request.method == "cancel": + if request.targetId is None: + raise OmpProtocolError( + "OMP UI cancellation is missing targetId" + ) + if turn is not None: + self._cancel_callback( + turn, OmpCallbackKind.UI, request.targetId + ) + return + if request.method in _UI_NOTIFICATIONS: + return + if turn is None: + write = asyncio.create_task( + self._write_frame( + ui_response_frame(request.id, OmpUIResponse(cancelled=True)) + ) + ) + self._write_tasks.add(write) + write.add_done_callback(self._write_tasks.discard) + write.add_done_callback(self._consume_task) + return + self._start_callback( + turn, OmpCallbackKind.UI, request.id, self._handle_ui(turn, request) + ) + case event_type if event_type in _EVENT_TYPES: + event = OmpEvent.model_validate(payload) + if turn is not None: + if overflow := turn.observe(event, sequence): + self._record_failure(overflow) + else: + logger.debug("OMP RPC ignored an event without an active prompt") + case _: + logger.debug("OMP RPC ignored a frame with an unknown type") + + def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: + if response.id is None: + raise OmpProtocolError("OMP response is missing its command id") + pending = self._pending.get(response.id) + turn = self._turn + active_prompt = turn if turn is not None and response.id == turn.id else None + expected_command = pending.command if pending is not None else None + if expected_command is None and active_prompt is not None: + expected_command = OmpCommandName.PROMPT + if expected_command is not None and response.command != expected_command: + raise OmpProtocolError("OMP response command does not match its request") + if pending is None: + if not response.success and active_prompt is not None: + active_prompt.rejected = True + active_prompt.fail(command_error(response)) + return + logger.debug("OMP RPC ignored an unmatched command response") + return + if pending.future.done(): + if ( + not response.success + and self._turn is not None + and response.id == self._turn.id + ): + self._turn.rejected = True + self._turn.fail(command_error(response)) + return + if response.success: + pending.future.set_result(ReceivedResponse(response, sequence)) + else: + pending.future.set_result(command_error(response)) + + def _start_callback( + self, + turn: PromptRun, + kind: OmpCallbackKind, + request_id: str, + callback: Coroutine[object, object, None], + ) -> None: + if self._closing or turn.failure.done(): + callback.close() + return + key = (kind, request_id) + task = asyncio.create_task(callback) + if key in turn.callbacks: + task.cancel() + raise OmpProtocolError("OMP callback request id was reused") + turn.callbacks[key] = task + turn.callback_tasks.add(task) + task.add_done_callback(lambda done: self._callback_finished(turn, key, done)) + + def _callback_still_active( + self, turn: PromptRun, kind: OmpCallbackKind, request_id: str + ) -> bool: + return ( + not self._closing + and turn.callbacks.get((kind, request_id)) is asyncio.current_task() + ) + + @staticmethod + def _cancel_callback( + turn: PromptRun, kind: OmpCallbackKind, request_id: str + ) -> None: + task = turn.callbacks.pop((kind, request_id), None) + if task is not None: + task.cancel() + + @staticmethod + def _callback_finished( + turn: PromptRun, key: tuple[OmpCallbackKind, str], task: asyncio.Task[None] + ) -> None: + if turn.callbacks.get(key) is task: + turn.callbacks.pop(key) + turn.callback_tasks.discard(task) + if not task.cancelled() and (error := task.exception()) is not None: + turn.fail(error) + + async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: + result = await turn.on_tool(call) + if not self._callback_still_active(turn, OmpCallbackKind.TOOL, call.id): + return + await self._write_frame( + OmpHostToolResult.from_tool(call.id, result).model_dump(mode="json") + ) + + async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: + response = await turn.on_ui(request) + if not self._callback_still_active(turn, OmpCallbackKind.UI, request.id): + return + await self._write_frame(ui_response_frame(request.id, response)) + + async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: + assert process.stderr is not None + try: + while data := await process.stderr.read(_DRAIN_BYTES): + logger.debug("OMP stderr drained: %s bytes", len(data)) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("OMP stderr reader failed") + + def _fail_pending(self, error: OmpProcessError) -> None: + for pending in self._pending.values(): + if not pending.future.done(): + pending.future.set_result(error) + if self._ready is not None and not self._ready.done(): + self._ready.set_result(error) + if self._turn is not None: + self._turn.fail(error) + + def _apply_ready_limits(self, ready: OmpReady) -> None: + self._max_frame_bytes = ready.maxFrameBytes + self._frames.max_reassembled_bytes = ready.maxReassembledFrameBytes + + def _record_failure(self, error: OmpProcessError) -> None: + if self._failure is None: + self._failure = error + logger.error("OMP RPC transport failed: %s", error) + self._fail_pending(self._failure) + if not self._closing: + self._begin_close() + + def _begin_close(self) -> None: + self._closing = True + if self._windows_process is not None: + self._windows_process.closing = True + if self._close_task is None or ( + self._close_task.done() + and ( + self._close_task.cancelled() or self._close_task.exception() is not None + ) + ): + self._close_task = asyncio.create_task(self.close_process()) + self._close_task.add_done_callback(self._log_close_failure) + + @staticmethod + def _log_close_failure(task: asyncio.Task[None]) -> None: + if not task.cancelled() and (error := task.exception()) is not None: + logger.error("OMP process cleanup failed", exc_info=error) + + @staticmethod + def _consume_task(task: asyncio.Task[None]) -> None: + if not task.cancelled(): + task.exception() + + @staticmethod + async def _cancel_tasks( + tasks: Sequence[asyncio.Task[None]], + *, + timeout: float, + label: str = "callback", + ) -> None: + active = {task for task in tasks if not task.done()} + for task in active: + task.cancel() + if active: + _, pending = await asyncio.wait(active, timeout=timeout) + for task in pending: + task.cancel() + logger.warning("OMP %s did not honor cancellation", label) + for task in tasks: + if task.done(): + OmpRpcClient._consume_task(task) + else: + task.add_done_callback(OmpRpcClient._consume_task) + + @staticmethod + async def _drain(stream: asyncio.StreamReader) -> None: + while await stream.read(_DRAIN_BYTES): + pass + + async def _await_contained_exit(self, process: asyncio.subprocess.Process) -> None: + if not await await_returncode(process, self.process_exit_timeout_s): + raise OmpProcessError("OMP process did not exit after containment kill") + + async def _settle_spawn(self) -> asyncio.subprocess.Process | None: + spawning = self._spawn_task + if spawning is None: + return self._proc + try: + return await asyncio.shield(spawning) + except Exception: # noqa: BLE001 -- startup propagates the error; cleanup retains its child. + return self._proc + except asyncio.CancelledError: + if not spawning.cancelled(): + raise + return self._proc + + async def _close_process(self) -> None: + turn = self._turn + self._fail_pending(self._failure or OmpProcessError("OMP RPC client closed")) + if turn is not None: + workers = list(turn.callback_tasks) + if turn.event_worker is not None: + workers.append(turn.event_worker) + await self._cancel_tasks( + workers, timeout=self.callback_exit_timeout_s, label="callback" + ) + + owner = self._windows_process + scope = self._process_scope + owner_closed = False + if owner is not None and self._proc is None: + # Stop incomplete helper creation/readiness before awaiting its spawn task. + try: + await owner.close(self.process_exit_timeout_s) + owner_closed = True + finally: + self._proc = owner.process + process = await self._settle_spawn() + if process is None: + self._windows_process = None + self._process_scope = None + return + if scope is not None: + await asyncio.to_thread(scope.capture, process) + if ( + self._failure is None + and process.returncode is None + and self._reader_task is not None + and not self._reader_task.done() + ): + try: + async with asyncio.timeout(_ABORT_TIMEOUT_S): + await self._request_frame( + OmpCommandName.ABORT, {}, allow_closing=True + ) + except (TimeoutError, OmpRpcError, OmpProcessError) as error: + logger.warning("OMP graceful abort skipped: %s", error) + reader = self._reader_task + if reader is not None and not reader.done(): + reader.cancel() + drains: list[asyncio.Task[None]] = [] + if ( + self._stderr_task is None or self._stderr_task.done() + ) and process.stderr is not None: + drains.append(asyncio.create_task(self._drain(process.stderr))) + try: + if scope is not None: + # Ownership lookup can disappear before process exit is observable. + await asyncio.to_thread(scope.capture, process) + await self._cancel_tasks( + list(self._write_tasks), + timeout=self.callback_exit_timeout_s, + label="write", + ) + if not owner_closed: + # A drain blocked on a full pipe keeps this lock; close must + # not wait for it. Writers are already cancelled. + await close_stdin(process) + await await_returncode(process, self.process_exit_timeout_s) + await asyncio.to_thread(self.after_stdin_closed) + if owner is not None: + if not owner_closed: + await owner.close(self.process_exit_timeout_s) + if scope is not None: + await asyncio.to_thread( + scope.close, process, self.process_exit_timeout_s + ) + else: + assert scope is not None + await asyncio.to_thread(scope.capture, process) + await kill_process_tree(process) + await asyncio.to_thread(scope.signal) + await self._await_contained_exit(process) + await asyncio.to_thread( + scope.close, process, self.process_exit_timeout_s + ) + await self._await_contained_exit(process) + self._windows_process = None + self._process_scope = None + finally: + # A cancelled pipe read on Windows stays pending until the child is + # killed; drain stdout only after that reader finishes. + if reader is not None: + await self._cancel_tasks( + [reader], timeout=self.callback_exit_timeout_s, label="reader" + ) + if process.stdout is not None: + drains.append(asyncio.create_task(self._drain(process.stdout))) + if self._stderr_task is not None: + drains.append(self._stderr_task) + await self._cancel_tasks( + drains, timeout=self.callback_exit_timeout_s, label="drain" + ) diff --git a/src/band/integrations/omprpc/errors.py b/src/band/integrations/omprpc/errors.py new file mode 100644 index 000000000..6c5820128 --- /dev/null +++ b/src/band/integrations/omprpc/errors.py @@ -0,0 +1,43 @@ +"""Native OMP RPC failure types.""" + +from __future__ import annotations + +from band.integrations.ompwire import OmpResponse + + +class OmpRpcError(RuntimeError): + """An unsuccessful command response; the transport remains usable.""" + + def __init__(self, command: str, error: str, *, code: str | None = None) -> None: + self.command = command + self.error = error + self.code = code + super().__init__(f"OMP {command}: {error}") + + +class OmpProcessError(RuntimeError): + """The OMP process or its stdio transport is no longer usable.""" + + +BACKLOG_EXCEEDED = "OMP event callback backlog exceeded" +CALLBACKS_UNFINISHED = "OMP callbacks did not finish after settlement" + + +class OmpProtocolError(OmpProcessError): + """An invalid or incorrectly correlated native RPC frame.""" + + +class OmpPromptError(RuntimeError): + """The prompt or its final background continuation failed or was aborted.""" + + def __init__(self, status: str, message: str) -> None: + self.status = status + super().__init__(message) + + +def command_error(response: OmpResponse) -> OmpRpcError: + return OmpRpcError( + response.command, + response.error or "Command failed", + code=response.code, + ) diff --git a/src/band/integrations/omprpc/frames.py b/src/band/integrations/omprpc/frames.py new file mode 100644 index 000000000..dbcd71dad --- /dev/null +++ b/src/band/integrations/omprpc/frames.py @@ -0,0 +1,117 @@ +"""Native OMP JSONL frame reassembly and host replies.""" + +from __future__ import annotations + +import base64 +from dataclasses import dataclass + +from pydantic import JsonValue, TypeAdapter + +from band.integrations.omprpc.errors import OmpProtocolError +from band.integrations.ompwire import ( + OMP_CHUNK_BYTES, + OMP_REASSEMBLED_BYTES, + OmpChunk, + OmpFrame, + OmpFrameType, + OmpUIResponse, +) + +_PAYLOAD_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) + + +@dataclass +class PendingChunks: + chunk_id: str + count: int + byte_length: int + next_index: int + received_bytes: int + parts: list[bytes] + + +class ChunkAssembler: + """Reassemble `rpc_chunk` sequences into one logical frame.""" + + def __init__(self) -> None: + self.max_reassembled_bytes = OMP_REASSEMBLED_BYTES + self._chunks: PendingChunks | None = None + + @property + def pending(self) -> bool: + return self._chunks is not None + + def reset(self) -> None: + self._chunks = None + + def decode(self, payload: dict[str, JsonValue]) -> dict[str, JsonValue] | None: + frame = OmpFrame.model_validate(payload) + if frame.type != OmpFrameType.CHUNK: + if self._chunks is not None: + raise OmpProtocolError("OMP RPC protocol chunk sequence interrupted") + return payload + chunk = OmpChunk.model_validate(payload) + if chunk.byteLength > self.max_reassembled_bytes: + raise OmpProtocolError( + "OMP RPC protocol frame exceeds the negotiated limit" + ) + data = base64.b64decode(chunk.data, validate=True) + if ( + not data + or len(data) > OMP_CHUNK_BYTES + or base64.b64encode(data).decode("ascii") != chunk.data + ): + raise OmpProtocolError("Invalid OMP RPC protocol chunk data") + if self._chunks is None: + if chunk.index != 0: + raise OmpProtocolError( + "OMP RPC protocol chunks must start at index zero" + ) + self._chunks = PendingChunks( + chunk.chunkId, chunk.count, chunk.byteLength, 0, 0, [] + ) + pending = self._chunks + if ( + pending.chunk_id != chunk.chunkId + or pending.count != chunk.count + or pending.byte_length != chunk.byteLength + or pending.next_index != chunk.index + ): + raise OmpProtocolError("OMP RPC protocol chunk sequence mismatch") + pending.parts.append(data) + pending.received_bytes += len(data) + pending.next_index += 1 + if pending.received_bytes > pending.byte_length: + raise OmpProtocolError( + "OMP RPC protocol chunks exceed their declared length" + ) + if pending.next_index < pending.count: + return None + if pending.received_bytes != pending.byte_length: + raise OmpProtocolError("OMP RPC protocol chunk length mismatch") + self._chunks = None + decoded = _PAYLOAD_ADAPTER.validate_json(b"".join(pending.parts), strict=True) + if OmpFrame.model_validate(decoded).type == OmpFrameType.CHUNK: + raise OmpProtocolError("OMP RPC protocol chunks cannot nest") + return decoded + + +def payload_json(frame: dict[str, JsonValue]) -> str: + return _PAYLOAD_ADAPTER.dump_json(frame).decode("utf-8") + + +def parse_payload_line(line: bytes) -> dict[str, JsonValue]: + return _PAYLOAD_ADAPTER.validate_json(line, strict=True) + + +def ui_response_frame(request_id: str, response: OmpUIResponse) -> dict[str, JsonValue]: + answer = response.model_dump( + mode="json", + exclude_none=True, + exclude=None if response.cancelled else {"cancelled"}, + ) + return { + "type": OmpFrameType.UI_RESPONSE, + "id": request_id, + **answer, + } diff --git a/src/band/integrations/omprpc/prompt.py b/src/band/integrations/omprpc/prompt.py new file mode 100644 index 000000000..88ab9590f --- /dev/null +++ b/src/band/integrations/omprpc/prompt.py @@ -0,0 +1,87 @@ +"""In-flight native OMP prompt state.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from enum import StrEnum + +from band.integrations.omprpc.errors import BACKLOG_EXCEEDED, OmpProcessError +from band.integrations.ompwire import ( + OMP_ASSISTANT_ROLE, + OmpEvent, + OmpEventType, + OmpHostToolCall, + OmpMessage, + OmpPromptResult, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +EVENT_BACKLOG_LIMIT = 256 + + +class OmpCallbackKind(StrEnum): + TOOL = "tool" + UI = "ui" + + +class PromptRun: + def __init__( + self, + request_id: str, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + loop = asyncio.get_running_loop() + self.id = request_id + self.on_event = on_event + self.on_tool = on_tool + self.on_ui = on_ui + self.events: asyncio.Queue[OmpEvent] = asyncio.Queue( + maxsize=EVENT_BACKLOG_LIMIT + ) + self.failure: asyncio.Future[BaseException] = loop.create_future() + self.result: asyncio.Future[OmpPromptResult] = loop.create_future() + self.changed = asyncio.Event() + self.last_activity_sequence = 0 + self.settled_sequence = 0 + self.last_assistant: OmpMessage | None = None + self.rejected = False + self.callbacks: dict[tuple[OmpCallbackKind, str], asyncio.Task[None]] = {} + self.callback_tasks: set[asyncio.Task[None]] = set() + self.event_worker: asyncio.Task[None] | None = None + + def fail(self, error: BaseException) -> None: + if not self.failure.done(): + self.failure.set_result(error) + + def observe(self, event: OmpEvent, sequence: int) -> OmpProcessError | None: + if self.failure.done(): + return None + # Waiting for space would block the reader that services reentrant commands. + try: + self.events.put_nowait(event) + except asyncio.QueueFull: + error = OmpProcessError(BACKLOG_EXCEEDED) + self.fail(error) + return error + self.last_activity_sequence = sequence + match event.type: + case OmpEventType.MESSAGE_END if ( + event.message is not None and event.message.role == OMP_ASSISTANT_ROLE + ): + self.last_assistant = event.message + case OmpEventType.AGENT_END: + self.last_assistant = next( + ( + message + for message in reversed(event.messages) + if message.role == OMP_ASSISTANT_ROLE + ), + self.last_assistant, + ) + self.changed.set() + return None diff --git a/src/band/integrations/omprpc/scope.py b/src/band/integrations/omprpc/scope.py new file mode 100644 index 000000000..4873e50ef --- /dev/null +++ b/src/band/integrations/omprpc/scope.py @@ -0,0 +1,111 @@ +"""Best-effort custody of an OMP RPC process tree.""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections.abc import Mapping +from uuid import uuid4 + +import psutil + +from band.integrations.codex.winspawn import session_pid_reused +from band.integrations.omprpc.errors import OmpProcessError + +_PROCESS_SCOPE_VALUE = "1" +# pid 0 is the kernel idle process; scanning it is wasted work. +KERNEL_IDLE_PID = 0 + + +class OmpProcessScope: + """Retain observable descendants for best-effort cleanup. + + The private PI_ marker survives native helper filters. Explicit environment + replacement and restricted process inspection can hide detached children. + Windows job membership is the primary containment; this marker catches + descendants the job does not observe. The job does not allow breakaway or + nest other jobs — a second job is not a substitute for this scan unless + close() still reports live members. + """ + + def __init__(self) -> None: + self.key = f"PI_BAND_PROCESS_{uuid4().hex}" + self.processes: set[psutil.Process] = set() + + def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: + return { + **(os.environ if base is None else base), + self.key: _PROCESS_SCOPE_VALUE, + } + + def capture(self, root: asyncio.subprocess.Process) -> None: + getpgid = getattr(os, "getpgid", None) + group_visible = getpgid is not None and not session_pid_reused(root) + group_members: list[psutil.Process] = [] + for pid in psutil.pids(): + # asyncio's child watcher remains the sole waiter for the root. + if pid == KERNEL_IDLE_PID or pid == root.pid: + continue + try: + member = psutil.Process(pid) + created = member.create_time() + marked = False + try: + marked = member.environ().get(self.key) == _PROCESS_SCOPE_VALUE + except ( + psutil.NoSuchProcess, + psutil.AccessDenied, + PermissionError, + SystemError, + ): + pass + in_group = False + if group_visible and getpgid is not None: + try: + in_group = getpgid(pid) == root.pid + except (ProcessLookupError, PermissionError): + pass + if not marked and not in_group: + continue + current = psutil.Process(pid) + if current.create_time() != created or current != member: + continue + if marked: + self.processes.add(member) + elif in_group: + group_members.append(member) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + # A reaped leader's id can be reused while membership is enumerated. + if group_visible and not session_pid_reused(root): + self.processes.update(group_members) + + @staticmethod + def _is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + def signal(self) -> None: + for process in self.processes: + try: + process.kill() + except psutil.NoSuchProcess: + pass + + def close(self, root: asyncio.subprocess.Process, timeout: float) -> None: + deadline = time.monotonic() + timeout + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + while alive: + self.signal() + psutil.wait_procs(alive, timeout=max(0.0, deadline - time.monotonic())) + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + if alive and time.monotonic() >= deadline: + raise OmpProcessError( + "OMP owned session members survived cleanup: " + f"{sorted(process.pid for process in alive)}" + ) diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py new file mode 100644 index 000000000..a1c14f4eb --- /dev/null +++ b/src/band/integrations/ompwire.py @@ -0,0 +1,376 @@ +"""Typed projections of OMP's native JSONL RPC frames.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Final, Literal, Self, cast + +from pydantic import ( + AliasChoices, + BaseModel, + ConfigDict, + Field, + JsonValue, + RootModel, + model_validator, +) + +OMP_FRAME_BYTES = 1024 * 1024 +OMP_REASSEMBLED_BYTES = 64 * 1024 * 1024 +OMP_CHUNK_BYTES = 256 * 1024 +OMP_NATIVE_PROTOCOL_VERSION: Final = 2 +OMP_ASSISTANT_ROLE = "assistant" + + +class OmpCommandName(StrEnum): + """Name SDK-issued commands without restricting arbitrary native requests.""" + + NEGOTIATE_PROTOCOL = "negotiate_protocol" + PROMPT = "prompt" + GET_STATE = "get_state" + GET_AVAILABLE_COMMANDS = "get_available_commands" + ABORT = "abort" + + +class OmpPromptStatus(StrEnum): + COMPLETED = "completed" + ERROR = "error" + ABORTED = "aborted" + + +class OmpFrameType(StrEnum): + READY = "ready" + RESPONSE = "response" + PROMPT_RESULT = "prompt_result" + SESSION_SETTLED = "session_settled" + HOST_TOOL_CALL = "host_tool_call" + HOST_TOOL_CANCEL = "host_tool_cancel" + HOST_TOOL_RESULT = "host_tool_result" + UI_REQUEST = "extension_ui_request" + UI_RESPONSE = "extension_ui_response" + CHUNK = "rpc_chunk" + FRAME_ERROR = "rpc_frame_error" + + +class OmpEventType(StrEnum): + AGENT_START = "agent_start" + AGENT_END = "agent_end" + TURN_START = "turn_start" + TURN_END = "turn_end" + MESSAGE_START = "message_start" + MESSAGE_UPDATE = "message_update" + MESSAGE_END = "message_end" + TOOL_START = "tool_execution_start" + TOOL_UPDATE = "tool_execution_update" + TOOL_STREAM = "tool_stream_update" + TOOL_END = "tool_execution_end" + + +class OmpWireModel(BaseModel): + model_config = ConfigDict( + frozen=True, extra="ignore", strict=True, allow_inf_nan=False + ) + + +class OmpContent(OmpWireModel): + """Transcript content, including text, thoughts, images, and tool calls.""" + + type: str + text: str | None = None + thinking: str | None = None + data: str | None = None + mimeType: str | None = None + id: str | None = None + name: str | None = None + arguments: dict[str, JsonValue] | None = None + intent: str | None = None + + +class OmpUsageCost(OmpWireModel): + input: float = 0 + output: float = 0 + cacheRead: float = 0 + cacheWrite: float = 0 + total: float = 0 + + +class OmpUsage(OmpWireModel): + input: int = 0 + output: int = 0 + cacheRead: int = 0 + cacheWrite: int = 0 + totalTokens: int = 0 + contextTokens: int | None = None + reasoningTokens: int | None = None + cost: OmpUsageCost | None = None + + +class OmpMessage(OmpWireModel): + role: str + content: str | list[OmpContent] = Field(default_factory=list) + stopReason: str | None = None + errorMessage: str | None = None + usage: OmpUsage | None = None + model: str | None = None + provider: str | None = None + api: str | None = None + timestamp: float | None = None + toolCallId: str | None = None + toolName: str | None = None + isError: bool | None = None + + +class OmpToolResult(OmpWireModel): + content: list[OmpContent] + details: JsonValue = None + is_error: bool = Field( + default=False, validation_alias=AliasChoices("is_error", "isError") + ) + + @model_validator(mode="after") + def validate_content(self) -> OmpToolResult: + for content in self.content: + match content.type: + case "text" if content.text is not None: + continue + case "image" if ( + content.data is not None and content.mimeType is not None + ): + continue + case _: + raise ValueError("OMP tool content must be text or an inline image") + return self + + +class OmpAssistantEvent(OmpWireModel): + type: str + delta: str | None = None + contentIndex: int | None = None + content: JsonValue = None + + +class OmpEvent(OmpWireModel): + type: str + message: OmpMessage | None = None + messages: list[OmpMessage] = Field(default_factory=list) + messageId: str | None = None + assistantMessageEvent: OmpAssistantEvent | None = None + toolCallId: str | None = None + toolName: str | None = None + args: JsonValue = None + result: OmpToolResult | None = None + partialResult: JsonValue = None + update: JsonValue = None + isError: bool | None = None + intent: str | None = None + isTerminal: bool | None = None + yielded: bool | None = None + awaitingAsyncWork: bool | None = None + + @model_validator(mode="after") + def validate_event(self) -> OmpEvent: + match self.type: + case ( + OmpEventType.MESSAGE_START + | OmpEventType.MESSAGE_UPDATE + | OmpEventType.MESSAGE_END + | OmpEventType.TURN_END + ): + if self.message is None: + raise ValueError("OMP message event is missing its message") + case OmpEventType.AGENT_END: + if "messages" not in self.model_fields_set: + raise ValueError("OMP agent_end is missing its messages") + case ( + OmpEventType.TOOL_START + | OmpEventType.TOOL_UPDATE + | OmpEventType.TOOL_STREAM + | OmpEventType.TOOL_END + ): + if self.toolCallId is None or self.toolName is None: + raise ValueError("OMP tool event is missing its call identity") + return self + + +class OmpModel(OmpWireModel): + id: str + provider: str + name: str | None = None + reasoning: bool | None = None + contextWindow: int | None = None + + +class OmpState(OmpWireModel): + sessionId: str + sessionFile: str | None = None + isSettled: bool + model: OmpModel | None = None + thinkingLevel: str | None = None + isStreaming: bool = False + hasPendingAsyncWork: bool = False + queuedMessageCount: int = 0 + + +class OmpHostToolCall(OmpWireModel): + type: Literal[OmpFrameType.HOST_TOOL_CALL] = OmpFrameType.HOST_TOOL_CALL + id: str + toolCallId: str + toolName: str + arguments: dict[str, JsonValue] + + +class OmpHostToolCancel(OmpWireModel): + type: Literal[OmpFrameType.HOST_TOOL_CANCEL] = OmpFrameType.HOST_TOOL_CANCEL + id: str + targetId: str + + +class OmpHostToolResult(OmpWireModel): + type: Literal[OmpFrameType.HOST_TOOL_RESULT] = OmpFrameType.HOST_TOOL_RESULT + id: str + result: dict[str, JsonValue] + isError: bool = False + + @classmethod + def from_tool(cls, call_id: str, tool: OmpToolResult) -> Self: + return cls( + id=call_id, + result=cast( + dict[str, JsonValue], + tool.model_dump(mode="json", exclude={"is_error"}, exclude_none=True), + ), + isError=tool.is_error, + ) + + +class OmpAskOption(OmpWireModel): + label: str + description: str | None = None + preview: str | None = None + + +class OmpAskQuestion(OmpWireModel): + """A structured native ask question, enabled by ``set_ask_dialog``.""" + + id: str + question: str + header: str | None = None + options: list[OmpAskOption] = Field(default_factory=list) + multi: bool = False + # The tool schema accepts any number and forwards the recommendation unchanged. + recommended: float | None = None + + +class OmpAskAnswer(OmpWireModel): + """One answer per question, in question order, using offered option labels.""" + + id: str + selectedOptions: list[str] = Field(default_factory=list) + customInput: str | None = None + + +class OmpUIRequest(OmpWireModel): + type: Literal[OmpFrameType.UI_REQUEST] = OmpFrameType.UI_REQUEST + id: str + method: str + title: str | None = None + options: list[str] = Field(default_factory=list) + message: str | None = None + targetId: str | None = None + questions: list[OmpAskQuestion] = Field(default_factory=list) + + +class OmpUIResponse(OmpWireModel): + """Answer or explicitly cancel one dialog.""" + + value: str | None = None + confirmed: bool | None = None + answers: list[OmpAskAnswer] | None = None + cancelled: bool = False + + @model_validator(mode="after") + def validate_answer(self) -> OmpUIResponse: + given = ( + self.value is not None, + self.confirmed is not None, + self.answers is not None, + self.cancelled, + ) + if sum(given) > 1: + raise ValueError("OMP UI response must contain only one answer") + return self + + +class OmpReady(OmpWireModel): + type: Literal[OmpFrameType.READY] + protocolVersion: Literal[1] + supportedProtocolVersions: list[int] + maxFrameBytes: int = Field(gt=0, le=OMP_FRAME_BYTES) + maxReassembledFrameBytes: int = Field(gt=0, le=OMP_REASSEMBLED_BYTES) + + +class OmpNegotiation(OmpWireModel): + protocolVersion: Literal[2] + + +class OmpChunk(OmpWireModel): + type: Literal[OmpFrameType.CHUNK] + chunkId: str = Field(min_length=1, max_length=128) + index: int = Field(ge=0) + count: int = Field(ge=2, le=OMP_REASSEMBLED_BYTES // OMP_CHUNK_BYTES) + byteLength: int = Field(ge=OMP_FRAME_BYTES, le=OMP_REASSEMBLED_BYTES) + data: str = Field(min_length=1) + + @model_validator(mode="after") + def validate_index(self) -> OmpChunk: + if self.index >= self.count: + raise ValueError("OMP chunk index exceeds its count") + return self + + +class OmpFrameError(OmpWireModel): + type: Literal[OmpFrameType.FRAME_ERROR] + error: str + + +class OmpResponse(OmpWireModel): + type: Literal[OmpFrameType.RESPONSE] + id: str | None = None + command: str + success: bool + data: JsonValue = None + error: str | None = None + code: str | None = None + + +class OmpPromptAcceptance(OmpWireModel): + agentInvoked: bool | None = None + + +class OmpPromptFailure(OmpWireModel): + message: str + retryable: bool + provider: str | None = None + model: str | None = None + httpStatus: int | None = None + + +class OmpPromptResult(OmpWireModel): + type: Literal[OmpFrameType.PROMPT_RESULT] + id: str | None = None + agentInvoked: bool + status: Literal[ + OmpPromptStatus.COMPLETED, + OmpPromptStatus.ERROR, + OmpPromptStatus.ABORTED, + ] + sessionSettled: bool + error: OmpPromptFailure | None = None + + +class OmpFrame(OmpWireModel): + type: str + + +class OmpCommand(RootModel[dict[str, JsonValue]]): + model_config = ConfigDict(strict=True, allow_inf_nan=False) diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py new file mode 100644 index 000000000..f59b1c758 --- /dev/null +++ b/tests/converters/test_omp.py @@ -0,0 +1,101 @@ +"""Tests for OmpHistoryConverter.""" + +from __future__ import annotations + +from band.converters.omp import OmpHistoryConverter, OmpMetadataKey +from band.core.types import MessageType + + +class TestOmpHistoryConverter: + def test_convert_empty_history(self) -> None: + state = OmpHistoryConverter().convert([]) + assert state.room_to_session == {} + assert state.replay_messages == [] + + def test_latest_session_file_wins_per_room(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": MessageType.TASK, + "metadata": { + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/old.jsonl", + }, + }, + { + "message_type": MessageType.TASK, + "metadata": { + OmpMetadataKey.ROOM: "room-2", + OmpMetadataKey.SESSION: "/sessions/other.jsonl", + }, + }, + { + "message_type": MessageType.TASK, + "metadata": { + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/new.jsonl", + }, + }, + ] + ) + assert state.room_to_session == { + "room-1": "/sessions/new.jsonl", + "room-2": "/sessions/other.jsonl", + } + + def test_unusable_metadata_is_ignored(self) -> None: + state = OmpHistoryConverter().convert( + [ + {"message_type": MessageType.TASK, "metadata": "not a mapping"}, + { + "message_type": MessageType.TASK, + "metadata": {OmpMetadataKey.ROOM: "room-1"}, + }, + { + "message_type": MessageType.TASK, + "metadata": { + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "", + }, + }, + { + "message_type": MessageType.TASK, + "metadata": { + OmpMetadataKey.ROOM: 7, + OmpMetadataKey.SESSION: "/sessions/x.jsonl", + }, + }, + ] + ) + assert state.room_to_session == {} + + def test_session_metadata_on_non_task_messages_is_ignored(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": MessageType.TEXT, + "metadata": { + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/injected.jsonl", + }, + "content": "hello", + "sender_name": "Alice", + "sender_type": "User", + } + ] + ) + assert state.room_to_session == {} + assert state.replay_messages == ["[Alice]: hello"] + + def test_room_transcript_is_replayed(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": MessageType.TEXT, + "content": "hello", + "sender_name": "Alice", + "sender_type": "User", + } + ] + ) + assert state.replay_messages == ["[Alice]: hello"] diff --git a/tests/framework_configs/converters.py b/tests/framework_configs/converters.py index 98b85b956..b930207df 100644 --- a/tests/framework_configs/converters.py +++ b/tests/framework_configs/converters.py @@ -319,9 +319,9 @@ def _build_gemini_config() -> ConverterConfig: # parsing and helpers are internal utility modules (shared helpers, not converters). # a2a / a2a_gateway use the A2A protocol which has a different message schema. # acp_client / acp_server use ACP protocol session updates, not standard convert(). -# codex, letta, and opencode are metadata-only converters that extract session state -# from task event metadata rather than converting message history. They don't implement -# the standard convert() -> framework-format contract that conformance tests validate. +# codex, omp, letta, and opencode extract session state from task-event metadata +# rather than converting message history. They don't implement the standard +# convert() -> framework-format contract that conformance tests validate. # crewai_flow is a metadata-only converter for orchestration state reconstructed from # task events; same exception as codex/letta/opencode. # slack is a metadata-only converter too: it recovers the Slack thread binding from the @@ -335,6 +335,7 @@ def _build_gemini_config() -> ConverterConfig: "acp_client", "acp_server", "codex", + "omp", "crewai_flow", "letta", "opencode", diff --git a/tests/integrations/codex/conftest.py b/tests/integrations/codex/conftest.py index 5d3fc862c..31eef2ca0 100644 --- a/tests/integrations/codex/conftest.py +++ b/tests/integrations/codex/conftest.py @@ -12,17 +12,30 @@ import pytest_asyncio from band.integrations.codex import stdio_client, winspawn +from tests.integrations.codex.job import JobKernel from tests.integrations.codex.lifecycle import WindowsLifecycle from tests.integrations.codex.lookup import WindowsLookup +@pytest.fixture +def job_kernel(monkeypatch: pytest.MonkeyPatch) -> JobKernel: + kernel = JobKernel() + monkeypatch.setattr(winspawn, "_kernel32", lambda: kernel) + monkeypatch.setattr(winspawn, "_last_error", lambda: kernel.last_error) + monkeypatch.setattr( + winspawn, "_win_os_error", lambda code: OSError(code, "native Job failure") + ) + return kernel + + @pytest.fixture def windows_lookup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> WindowsLookup: peer = WindowsLookup(tmp_path) monkeypatch.chdir(peer.current) monkeypatch.setenv("PATH", "") - monkeypatch.setattr(stdio_client, "os", SimpleNamespace(name="nt", path=os.path)) - monkeypatch.setattr(winspawn, "os", SimpleNamespace(name="nt", pathsep=os.pathsep)) + monkeypatch.setattr( + winspawn, "os", SimpleNamespace(name="nt", path=os.path, pathsep=os.pathsep) + ) monkeypatch.setattr(winspawn, "_kernel32", lambda: peer.kernel) return peer @@ -40,9 +53,6 @@ async def windows_lifecycle( ) monkeypatch.setattr(asyncio, "create_subprocess_exec", peer.launch) monkeypatch.setattr(winspawn, "_read_status", peer.read_status) - monkeypatch.setattr( - peer.client, "_windows_command", lambda _job, _handle: ["unused"] - ) try: yield peer finally: diff --git a/tests/integrations/codex/job.py b/tests/integrations/codex/job.py new file mode 100644 index 000000000..525a02868 --- /dev/null +++ b/tests/integrations/codex/job.py @@ -0,0 +1,166 @@ +"""Control Job accounting independently of native process-exit signals.""" + +from __future__ import annotations + +import ctypes +from ctypes import wintypes +from typing import Any + +from band.integrations.codex.winspawn import ( + BasicAccountingInformation, + ExtendedLimitInformation, + ProcessIdList, +) + +JOB_HANDLE = 1 +FIRST_MEMBER_PID = 21 +SECOND_MEMBER_PID = 22 +PROCESS_HANDLES = {FIRST_MEMBER_PID: 121, SECOND_MEMBER_PID: 122} +KILL_ON_JOB_CLOSE = 0x2000 +UNRELATED_HANDLE = 999 + + +def handle_value(handle: ctypes.c_void_p | int) -> int: + if isinstance(handle, int): + return handle + assert handle.value is not None + return handle.value + + +class JobKernel: + def __init__(self) -> None: + self.pids = list(PROCESS_HANDLES) + self.signaled: set[int] = set() + self.opened: list[int] = [] + self.closed: list[int] = [] + self.terminated: list[int] = [] + self.termination_attempts: list[int] = [] + self.close_failures: dict[int, int] = {} + self.termination_failures = 0 + self.wait_fails = False + self.last_error = 0 + self.admission_sealed = False + self.fork_during_open = False + self.late_fork_rejected = False + self.recycled_pid: int | None = None + self.missing_pid: int | None = None + self.inaccessible_pid: int | None = None + self.limits = ExtendedLimitInformation() + self.limits.BasicLimitInformation.LimitFlags = KILL_ON_JOB_CLOSE + + def QueryInformationJobObject( + self, + job: ctypes.c_void_p | int, + information_class: int, + information: Any, + size: int, + returned: object, + ) -> bool: + assert handle_value(job) == JOB_HANDLE + match information_class: + case 1: + accounting = ctypes.cast( + information, ctypes.POINTER(BasicAccountingInformation) + ).contents + accounting.ActiveProcesses = 0 if self.terminated else len(self.pids) + return True + case 3: + header = ctypes.cast( + information, ctypes.POINTER(ProcessIdList) + ).contents + pids = [] if self.terminated else self.pids + capacity = (size - ProcessIdList.ProcessIdList.offset) // ctypes.sizeof( + ctypes.c_size_t + ) + header.NumberOfAssignedProcesses = len(pids) + header.NumberOfProcessIdsInList = min(len(pids), capacity) + address = ctypes.addressof(header) + ProcessIdList.ProcessIdList.offset + buffer = (ctypes.c_size_t * capacity).from_address(address) + for index, pid in enumerate(pids[:capacity]): + buffer[index] = pid + if len(pids) > capacity: + self.last_error = 234 + return False + return True + case 9: + assert size == ctypes.sizeof(self.limits) + ctypes.memmove(information, ctypes.byref(self.limits), size) + return True + case _: + raise AssertionError(f"Unexpected Job query class: {information_class}") + + def SetInformationJobObject( + self, + job: ctypes.c_void_p | int, + information_class: int, + information: Any, + size: int, + ) -> bool: + assert handle_value(job) == JOB_HANDLE + assert information_class == 9 + assert size == ctypes.sizeof(self.limits) + ctypes.memmove(ctypes.byref(self.limits), information, size) + basic = self.limits.BasicLimitInformation + self.admission_sealed = ( + bool(basic.LimitFlags & 8) and basic.ActiveProcessLimit == 0 + ) + return True + + def OpenProcess(self, rights: int, inherit: bool, pid: int) -> int | None: + assert rights & 0x00100000, "process handles must support waiting" + assert rights & 0x1000, "process handles must support Job membership queries" + assert not inherit + if self.fork_during_open: + self.fork_during_open = False + self.late_fork_rejected = self.admission_sealed + if not self.admission_sealed: + self.pids.append(23) + if pid == self.missing_pid: + self.last_error = 87 + return None + if pid == self.inaccessible_pid: + self.last_error = 5 + return None + handle = UNRELATED_HANDLE if pid == self.recycled_pid else PROCESS_HANDLES[pid] + self.opened.append(handle) + return handle + + def IsProcessInJob( + self, process: ctypes.c_void_p | int, job: ctypes.c_void_p | int, result: Any + ) -> bool: + assert handle_value(job) == JOB_HANDLE + ctypes.cast(result, ctypes.POINTER(wintypes.BOOL)).contents.value = ( + handle_value(process) != UNRELATED_HANDLE + ) + return True + + def TerminateJobObject(self, job: ctypes.c_void_p | int, code: int) -> bool: + assert code == 1 + self.termination_attempts.append(handle_value(job)) + if self.termination_failures: + self.termination_failures -= 1 + self.last_error = 5 + return False + self.terminated.append(handle_value(job)) + return True + + def WaitForSingleObject(self, process: ctypes.c_void_p | int, timeout: int) -> int: + assert timeout == 0, "native probes must not block the event loop" + if self.wait_fails: + self.last_error = 6 + return 0xFFFFFFFF + return 0 if handle_value(process) in self.signaled else 258 + + def CloseHandle(self, handle: ctypes.c_void_p | int) -> bool: + value = handle_value(handle) + remaining = self.close_failures.get(value, 0) + if remaining: + self.close_failures[value] = remaining - 1 + self.last_error = 5 + return False + self.closed.append(value) + return True + + def assert_handles_released(self) -> None: + """Retirement releases every handle acquired across failed attempts.""" + assert sorted(self.closed) == sorted([JOB_HANDLE, *self.opened]) diff --git a/tests/integrations/codex/lifecycle.py b/tests/integrations/codex/lifecycle.py index c6f2bb2ef..acd0fad2f 100644 --- a/tests/integrations/codex/lifecycle.py +++ b/tests/integrations/codex/lifecycle.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import sys from collections.abc import Mapping from typing import cast @@ -65,7 +66,7 @@ class WindowsLifecycle: def __init__(self) -> None: self.helper = HelperPeer() self.job = JobPeer(self.helper) - self.client = CodexStdioClient(command=["unused"]) + self.client = CodexStdioClient(command=[sys.executable]) self.creation_started = asyncio.Event() self.creation_allowed = asyncio.Event() self.creation_allowed.set() diff --git a/tests/integrations/codex/lookup.py b/tests/integrations/codex/lookup.py index 1aebc9ac1..b2820f664 100644 --- a/tests/integrations/codex/lookup.py +++ b/tests/integrations/codex/lookup.py @@ -9,6 +9,7 @@ from types import SimpleNamespace from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job async def launch_directly(client: CodexStdioClient) -> None: @@ -30,7 +31,7 @@ async def launch_directly(client: CodexStdioClient) -> None: def selected_executable(name: str) -> Path: - command = CodexStdioClient(command=[name])._windows_command("test-job", 1) + command = command_in_job([name], "test-job", 1) return Path(command[command.index("--") + 1]) diff --git a/tests/integrations/codex/test_bootstrap.py b/tests/integrations/codex/test_bootstrap.py index 41e47c799..0ad48891e 100644 --- a/tests/integrations/codex/test_bootstrap.py +++ b/tests/integrations/codex/test_bootstrap.py @@ -10,6 +10,7 @@ import pytest from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job class TestWindowsBootstrap: @@ -38,7 +39,7 @@ def test_room_pythonpath_cannot_supply_helper_code( cwd=str(workdir), env={"PYTHONPATH": pythonpath} if configured else None, ) - command = client._windows_command("job", 1) + command = command_in_job(client.command, "job", 1) completed = subprocess.run( command[:4], cwd=workdir, @@ -73,7 +74,7 @@ def test_child_path_does_not_change_parent_executable_selection( client = CodexStdioClient( command=[executable.name, "-c", "pass"], env={"PATH": str(child_path)} ) - command = client._windows_command("job", 1) + command = command_in_job(client.command, "job", 1) selected = command[command.index("--") + 1] assert Path(selected).samefile(executable) assert client.env is not None and client.env["PATH"] == str(child_path) diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index b382d7644..ae555a0a9 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import ctypes import logging import os import shutil @@ -19,13 +20,15 @@ import pytest_asyncio from band.integrations.codex import CodexStdioClient, stdio_client -from band.integrations.codex.stdio_client import HELPER_PATH, kill_process_tree from band.integrations.codex.winspawn import ( + HELPER_PATH, Job, JobNotDrained, SpawnFailure, _read_status, _start_and_wait, + command_in_job, + kill_process_tree, ) from tests.integrations.codex.cleanup import cleanup_snapshot from tests.paths import ROOM_PEER @@ -223,8 +226,7 @@ async def test_forced_close_stops_what_a_launcher_shim_spawned( async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( self, tmp_path: Path ) -> None: - """A launcher that exits on stdin EOF would otherwise leave its child - holding the workdir.""" + """POSIX launchers observe stdin EOF; every platform reaps their worker.""" workdir = tmp_path / "room" workdir.mkdir() client = CodexStdioClient( @@ -237,7 +239,8 @@ async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( await asyncio.wait_for(wait_for_heartbeat(heartbeat), timeout=CLOSE_BOUND_S) await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) await assert_heartbeat_stopped(heartbeat) - assert not (workdir / STDIN_EOF_MARKER).exists() + if sys.platform != "win32": + assert (workdir / STDIN_EOF_MARKER).exists() shutil.rmtree(workdir) finally: await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) @@ -311,27 +314,19 @@ def deny(_name: str) -> object: async def test_idle_wait_preserves_ownership_until_explicit_release( self, active: int | None, drained: bool ) -> None: - job = Job.__new__(Job) - job._handle = object() # type: ignore[assignment] - job.name = "test" - closed = False + handle = ctypes.c_void_p(1) + job = Job(handle, "test") def active_processes() -> int | None: return active - def close() -> None: - nonlocal closed - closed = True - job._handle = None - job.active_processes = active_processes # type: ignore[method-assign] - job.close = close # type: ignore[method-assign] if drained: await job.wait_until_idle(0) else: with pytest.raises(JobNotDrained): await job.wait_until_idle(0) - assert not closed + assert job._handle is handle def test_windows_helper_does_not_import_a_room_package( self, tmp_path: Path @@ -343,7 +338,7 @@ def test_windows_helper_does_not_import_a_room_package( command=[sys.executable, "-c", "pass"], cwd=str(workdir), ) - argv = client._windows_command("job", 1) + argv = command_in_job(client.command, "job", 1) completed = subprocess.run( argv[:4], cwd=workdir, @@ -379,11 +374,9 @@ def test_windows_command_does_not_resolve_a_binary_in_the_room( cwd=str(workdir), ) with pytest.raises(FileNotFoundError): - client._windows_command("job", 1) + command_in_job(client.command, "job", 1) named = Path(sys.executable).name - resolved = CodexStdioClient(command=[named, "-c", "pass"])._windows_command( - "job", 1 - ) + resolved = command_in_job([named, "-c", "pass"], "job", 1) assert Path(resolved[resolved.index("--") + 1]).is_absolute() @pytest.mark.skipif(os.name == "nt", reason="session-group teardown is Unix") diff --git a/tests/integrations/codex/test_job.py b/tests/integrations/codex/test_job.py index c801264cf..7f5053fbc 100644 --- a/tests/integrations/codex/test_job.py +++ b/tests/integrations/codex/test_job.py @@ -2,12 +2,22 @@ from __future__ import annotations +import asyncio import ctypes from types import SimpleNamespace import pytest from band.integrations.codex import winspawn +from tests.integrations.codex.job import ( + FIRST_MEMBER_PID, + JOB_HANDLE, + KILL_ON_JOB_CLOSE, + PROCESS_HANDLES, + SECOND_MEMBER_PID, + UNRELATED_HANDLE, + JobKernel, +) class TestJobHandleOwnership: @@ -32,25 +42,168 @@ def close_handle(current: ctypes.c_void_p) -> bool: job.close() assert attempted == [handle, handle] - def test_failed_termination_is_reported_without_releasing_the_job( - self, monkeypatch: pytest.MonkeyPatch + async def test_failed_termination_is_reported_without_releasing_the_job( + self, job_kernel: JobKernel ) -> None: - handle = ctypes.c_void_p(1) + handle = ctypes.c_void_p(JOB_HANDLE) job = winspawn.Job(handle, "test") - terminated: list[ctypes.c_void_p] = [] + job_kernel.termination_failures = 1 + with pytest.raises(OSError): + job.terminate() + assert not job_kernel.closed + job.terminate() + assert job_kernel.termination_attempts == [JOB_HANDLE, JOB_HANDLE] + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] - def terminate(current: ctypes.c_void_p, _code: int) -> bool: - terminated.append(current) - return len(terminated) > 1 + async def test_empty_accounting_does_not_release_pending_kernel_processes( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + assert job.active_processes() == 0 + with pytest.raises(winspawn.JobNotDrained): + await job.wait_until_idle(0) + assert not job_kernel.closed - monkeypatch.setattr( - winspawn, - "_kernel32", - lambda: SimpleNamespace(TerminateJobObject=terminate), - ) - monkeypatch.setattr(winspawn, "_last_error", lambda: 5) - monkeypatch.setattr(winspawn, "_win_os_error", lambda code: OSError(code)) + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + def test_admission_is_sealed_before_a_member_can_fork_during_capture( + self, job_kernel: JobKernel + ) -> None: + job_kernel.fork_during_open = True + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + assert job_kernel.late_fork_rejected + assert job_kernel.limits.BasicLimitInformation.LimitFlags & KILL_ON_JOB_CLOSE + + async def test_recycled_pid_is_not_waited_as_an_owned_process( + self, job_kernel: JobKernel + ) -> None: + job_kernel.recycled_pid = FIRST_MEMBER_PID + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [ + JOB_HANDLE, + PROCESS_HANDLES[SECOND_MEMBER_PID], + UNRELATED_HANDLE, + ] + + async def test_process_that_exits_during_handle_capture_does_not_block_cleanup( + self, job_kernel: JobKernel + ) -> None: + job_kernel.missing_pid = FIRST_MEMBER_PID + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [ + JOB_HANDLE, + PROCESS_HANDLES[SECOND_MEMBER_PID], + ] + + async def test_native_wait_failure_retains_process_handles_for_retry( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.wait_fails = True + with pytest.raises(OSError): + await job.wait_until_idle(0) + assert not job_kernel.closed + job_kernel.wait_fails = False + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + @pytest.mark.parametrize("failing_pid", [FIRST_MEMBER_PID, SECOND_MEMBER_PID]) + async def test_failed_member_handle_close_keeps_the_job_owned_for_retry( + self, job_kernel: JobKernel, failing_pid: int + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job_kernel.close_failures[PROCESS_HANDLES[failing_pid]] = 1 + with pytest.raises(OSError): + job.close() + assert JOB_HANDLE not in job_kernel.closed + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + async def test_unrelated_identity_handle_close_failure_is_retained_for_retry( + self, job_kernel: JobKernel + ) -> None: + job_kernel.recycled_pid = FIRST_MEMBER_PID + job_kernel.close_failures[UNRELATED_HANDLE] = 1 + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") with pytest.raises(OSError): job.terminate() + assert not job_kernel.closed + assert not job_kernel.terminated job.terminate() - assert terminated == [handle, handle] + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + job_kernel.assert_handles_released() + + async def test_partial_capture_failure_retains_the_handles_already_opened( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job_kernel.inaccessible_pid = SECOND_MEMBER_PID + with pytest.raises(OSError): + job.terminate() + assert not job_kernel.closed + assert not job_kernel.terminated + job_kernel.inaccessible_pid = None + job.terminate() + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + job_kernel.assert_handles_released() + + async def test_cancelled_native_exit_wait_keeps_handles_for_explicit_retry( + self, job_kernel: JobKernel, monkeypatch: pytest.MonkeyPatch + ) -> None: + entered = asyncio.Event() + native_wait = job_kernel.WaitForSingleObject + + def observed_wait(handle: ctypes.c_void_p | int, timeout: int) -> int: + entered.set() + return native_wait(handle, timeout) + + monkeypatch.setattr(job_kernel, "WaitForSingleObject", observed_wait) + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + waiting = asyncio.create_task(job.wait_until_idle(30)) + observed = asyncio.create_task(entered.wait()) + try: + done, _pending = await asyncio.wait( + [waiting, observed], return_when=asyncio.FIRST_COMPLETED + ) + assert observed in done + assert not waiting.done() + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + assert not job_kernel.closed + job_kernel.signaled.update(PROCESS_HANDLES.values()) + job.terminate() + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + finally: + waiting.cancel() + observed.cancel() + await asyncio.gather(waiting, observed, return_exceptions=True) diff --git a/tests/integrations/codex/test_lookup.py b/tests/integrations/codex/test_lookup.py index 932560ad4..963a36156 100644 --- a/tests/integrations/codex/test_lookup.py +++ b/tests/integrations/codex/test_lookup.py @@ -7,7 +7,7 @@ import pytest -from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job from tests.integrations.codex.lookup import WindowsLookup, selected_executable @@ -16,9 +16,7 @@ def test_helper_uses_the_running_image_instead_of_a_venv_launcher( self, windows_lookup: WindowsLookup ) -> None: peer = windows_lookup - command = CodexStdioClient(command=[str(peer.image)])._windows_command( - "test-job", 1 - ) + command = command_in_job([str(peer.image)], "test-job", 1) assert Path(command[0]).samefile(peer.image) @pytest.mark.parametrize( diff --git a/tests/integrations/codex/test_stdio_client.py b/tests/integrations/codex/test_stdio_client.py index ff255cc58..332f425b1 100644 --- a/tests/integrations/codex/test_stdio_client.py +++ b/tests/integrations/codex/test_stdio_client.py @@ -408,3 +408,17 @@ async def test_stdio_client_eof_fails_pending_and_enqueues_disconnect( event = await client.recv_event(timeout_s=1.0) assert event.method == "transport/closed" await client.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="Windows closes through the job object") +async def test_close_lets_codex_exit_on_stdin_eof_before_killing_its_tree( + fake_codex_server_script: Path, +) -> None: + """A harness that exits once its stdin closes ends on its own terms; the + tree kill is the fallback for one that does not.""" + client = await _build_client(fake_codex_server_script, scenario="basic") + process = client._proc + assert process is not None + await client.close() + assert process.returncode == 0 diff --git a/tests/integrations/codex/test_windows.py b/tests/integrations/codex/test_windows.py index 278f651ce..4a8f8cde5 100644 --- a/tests/integrations/codex/test_windows.py +++ b/tests/integrations/codex/test_windows.py @@ -4,13 +4,15 @@ import os import sys +from collections.abc import Sequence from pathlib import Path import pytest from band.core.exceptions import BandConnectionError -from band.integrations.codex.stdio_client import HELPER_PATH, CodexStdioClient -from band.integrations.codex.winspawn import WindowsProcess +from band.integrations.codex import winspawn +from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import HELPER_PATH, WindowsProcess, command_in_job @pytest.mark.skipif(os.name != "nt", reason="requires native Windows job objects") @@ -57,12 +59,14 @@ async def test_rejected_spawn_releases_helper_and_tree_directory_handles( command=[sys.executable, "-c", "import sys; sys.stdin.buffer.read()"], cwd=str(workdir), ) - command_for = client._windows_command + command_for = command_in_job report = tmp_path / "helper-pid" owners: list[WindowsProcess] = [] - def rejected_command(job_name: str, status_handle: int) -> list[str]: - command = command_for(job_name, status_handle) + def rejected_command( + child_command: Sequence[str], job_name: str, status_handle: int + ) -> list[str]: + command = command_for(child_command, job_name, status_handle) assert client._windows_process is not None owners.append(client._windows_process) return [ @@ -74,7 +78,7 @@ def rejected_command(job_name: str, status_handle: int) -> list[str]: *command[4:], ] - monkeypatch.setattr(client, "_windows_command", rejected_command) + monkeypatch.setattr(winspawn, "command_in_job", rejected_command) try: with pytest.raises(BandConnectionError): await client.connect() diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py new file mode 100644 index 000000000..76c797ee7 --- /dev/null +++ b/tests/integrations/omppeer.py @@ -0,0 +1,906 @@ +"""Local OMP RPC peer with input-driven settlement and cancellation barriers.""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import signal +import socket +import subprocess +import sys +import threading +from typing import cast + +import psutil +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from band.integrations.ompwire import ( + OMP_CHUNK_BYTES, + OMP_FRAME_BYTES, + OMP_NATIVE_PROTOCOL_VERSION, + OMP_REASSEMBLED_BYTES, + OmpCommandName, + OmpPromptStatus, +) + +PROTOCOL_VERSION = 1 +COMPLETE = "complete" +YIELDED = "yielded" +BACKGROUND_COMPLETE = "background complete" +RESUMED_PROVIDER_FAILED = "Resumed provider failed" +FOLLOW_UP_PROMPT = "second" +CRASH_EXIT = 3 +FRAME_PAYLOAD_MARK = "omp-frame-payload-must-not-be-logged" +UI_APPROVE = "Approve" +UI_DENY = "Deny" +HOST_CANCELLED_ID = "host-cancelled" +UI_CANCELLED_ID = "ui-cancelled" +OVERSIZE_AFTER_READY = "oversize-after-ready" +READY_LIMIT_BYTES = 1024 +LATE_PROMPT_ERROR = "late-prompt-error" +ADMITTED_PROMPT_ERROR = "late-error" +SETTLED_CALLBACK_SCENARIO = "settled-during-callback" +SETTLED_CALLBACK_ID = "host-held" +LATE_CALLBACK_ID = "host-late" +STDIN_BACKPRESSURE = "stdin-backpressure" + + +def burst_message_id(number: int) -> str: + return f"burst-{number}" + + +class Frame(BaseModel): + model_config = ConfigDict(extra="allow") + __pydantic_extra__: dict[str, JsonValue] = Field(init=False) + + type: str + id: str | None = None + message: str | None = None + failed_command: str | None = None + result: dict[str, JsonValue] | None = None + isError: bool | None = None + value: str | None = None + confirmed: bool | None = None + cancelled: bool | None = None + count: int = 0 + + +class ProcessReceipt(BaseModel): + pid: int + parent_pid: int + created: float + + +def start_descendant( + scenario: str, + port: int, + stdio: str, + *, + detached: bool = False, + spawn_on_eof: bool = False, +) -> subprocess.Popen[bytes]: + return subprocess.Popen( + ( + sys.executable, + "-u", + __file__, + "--scenario", + scenario, + "--gate-port", + str(port), + *(("--detached-children",) if detached else ()), + *(("--spawn-on-eof",) if spawn_on_eof else ()), + ), + # A parent's pending pipe read can block Python's stdin initialization on Windows. + # Descendants are controlled by the gate socket, not the RPC input pipe. + stdin=subprocess.PIPE if spawn_on_eof else subprocess.DEVNULL, + stdout=None if stdio == "inherit" else subprocess.DEVNULL, + stderr=None if stdio == "inherit" else subprocess.DEVNULL, + start_new_session=detached and os.name != "nt", + ) + + +def hold_process_gate( + port: int, + *, + spawn_child: bool, + detached: bool = False, + spawn_on_eof: bool = False, +) -> None: + if hasattr(signal, "SIGTERM"): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + + def launch_child() -> subprocess.Popen[bytes] | None: + return ( + start_descendant("gated-child", port, "inherit", detached=detached) + if spawn_child + else None + ) + + child = None if spawn_on_eof else launch_child() + with socket.create_connection(("127.0.0.1", port), timeout=5) as gate: + gate.settimeout(None) + receipt = ProcessReceipt( + pid=os.getpid(), + parent_pid=os.getppid(), + created=psutil.Process().create_time(), + ) + gate.sendall((receipt.model_dump_json() + "\n").encode()) + if spawn_on_eof: + sys.stdin.buffer.read() + child = launch_child() + gate.recv(1) + if child is not None: + child.wait(timeout=5) + + +def send(frame: dict[str, JsonValue]) -> None: + payload = json.dumps(frame, separators=(",", ":")).encode("utf-8") + if ( + PROTOCOL_VERSION == OMP_NATIVE_PROTOCOL_VERSION + and len(payload) + 1 > OMP_FRAME_BYTES + ): + count = (len(payload) + OMP_CHUNK_BYTES - 1) // OMP_CHUNK_BYTES + for index in range(count): + chunk: dict[str, JsonValue] = { + "type": "rpc_chunk", + "chunkId": "large-frame", + "index": index, + "count": count, + "byteLength": len(payload), + "data": base64.b64encode( + payload[index * OMP_CHUNK_BYTES : (index + 1) * OMP_CHUNK_BYTES] + ).decode("ascii"), + } + sys.stdout.write(json.dumps(chunk, separators=(",", ":")) + "\n") + else: + sys.stdout.write(payload.decode("utf-8") + "\n") + sys.stdout.flush() + + +def assistant(text: str, *, stop: str = "stop") -> dict[str, JsonValue]: + message: dict[str, JsonValue] = { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Considering the result"}, + {"type": "text", "text": text}, + ], + "api": "test-api", + "provider": "test-provider", + "model": "test-model", + "stopReason": stop, + "timestamp": 1, + "usage": { + "input": 7, + "output": 3, + "cacheRead": 2, + "cacheWrite": 1, + "totalTokens": 13, + "cost": { + "input": 0.1, + "output": 0.2, + "cacheRead": 0.01, + "cacheWrite": 0.02, + "total": 0.33, + }, + }, + } + if stop == OmpPromptStatus.ERROR: + message["errorMessage"] = RESUMED_PROVIDER_FAILED + return message + + +class Peer: + def __init__(self, scenario: str) -> None: + self.scenario = scenario + self.settled = True + self.state_requested = False + self.prompt_id: str | None = None + self.prompt_count = 0 + self.frames: list[dict[str, JsonValue]] = [] + self.tool_results: set[str] = set() + self.ui_answered = False + self.held: Frame | None = None + + def respond(self, frame: Frame, data: dict[str, JsonValue] | None = None) -> None: + result: dict[str, JsonValue] = { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": True, + } + if data is not None: + result["data"] = data + send(result) + + def checkpoint(self, frame: Frame) -> None: + if ( + self.scenario == SETTLED_CALLBACK_SCENARIO + and self.prompt_id is not None + and not self.state_requested + ): + self.held = frame + return + self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) + + def end( + self, + text: str, + *, + stop: str = "stop", + pending: bool = False, + compacted: bool = False, + ) -> None: + message = assistant(text, stop=stop) + send({"type": "message_end", "messageId": text, "message": message}) + send( + { + "type": "agent_end", + "messages": [] if compacted else [message], + "isTerminal": not pending, + "yielded": True, + "awaitingAsyncWork": pending, + } + ) + + def result(self, *, settled: bool, status: str = "completed") -> None: + result: dict[str, JsonValue] = { + "type": "prompt_result", + "id": self.prompt_id, + "agentInvoked": True, + "status": status, + "sessionSettled": settled, + } + if status == "error": + result["error"] = {"message": "Initial provider failed", "retryable": False} + send(result) + + def settle(self) -> None: + self.settled = True + send({"type": "session_settled"}) + + def finish(self, frame: Frame) -> None: + self.respond(frame) + send({"type": "agent_start"}) + self.end(COMPLETE) + self.result(settled=True) + self.settle() + + def invalid_chunk(self) -> None: + chunk: dict[str, JsonValue] = { + "type": "rpc_chunk", + "chunkId": "invalid-frame", + "index": 0, + "count": 2, + "byteLength": OMP_FRAME_BYTES, + "data": "e30=", + } + match self.scenario: + case "chunk-overflow": + chunk["byteLength"] = OMP_REASSEMBLED_BYTES + 1 + send(chunk) + case "chunk-order": + chunk["index"] = 1 + send(chunk) + case "chunk-invalid-data": + chunk["data"] = "!" + send(chunk) + case "chunk-length": + send(chunk) + chunk["index"] = 1 + send(chunk) + case "chunk-eof": + send(chunk) + raise SystemExit(2) + case _: + send(chunk) + send({"type": "notice", "level": "info", "message": "Interrupted"}) + + def prompt(self, frame: Frame) -> None: + self.prompt_id = frame.id + self.prompt_count += 1 + self.settled = False + self.state_requested = False + if ( + self.prompt_count > 1 and self.scenario != "lifecycle" + ) or frame.message == FOLLOW_UP_PROMPT: + self.finish(frame) + return + match self.scenario: + case "rejected": + self.settled = True + send( + { + "type": "response", + "id": frame.id, + "command": OmpCommandName.PROMPT, + "success": False, + "error": "Prompt rejected", + "code": "rejected", + } + ) + return + case scenario if scenario == ADMITTED_PROMPT_ERROR: + self.respond(frame) + send( + { + "type": "response", + "id": frame.id, + "command": OmpCommandName.PROMPT, + "success": False, + "error": "Admitted prompt failed later", + "code": "post_admit", + } + ) + return + case "local": + self.settled = True + self.respond(frame, {"agentInvoked": False}) + return + case "bad-acceptance": + self.respond(frame, {"agentInvoked": "yes"}) + return + case "eof": + self.respond(frame) + send({"type": "agent_start"}) + raise SystemExit(2) + case "bad-json": + sys.stdout.write("{" + FRAME_PAYLOAD_MARK + "\n") + sys.stdout.flush() + return + case "bad-event": + self.respond(frame) + send({"type": "message_end", "message": {"role": 123}}) + return + case "wrong-command": + send( + { + "type": "response", + "id": frame.id, + "command": OmpCommandName.GET_STATE, + "success": True, + "data": {}, + } + ) + return + case "callbacks": + self.held = frame + send({"type": "agent_start"}) + for number in (1, 2): + send( + { + "type": "host_tool_call", + "id": f"host-{number}", + "toolCallId": f"call-{number}", + "toolName": "band_chat", + "arguments": {"number": number}, + } + ) + send( + { + "type": "extension_ui_request", + "id": "approval-1", + "method": "select", + "title": "Allow command?", + "options": [UI_APPROVE, UI_DENY], + } + ) + return + case "stale": + send({"type": "session_settled"}) + send( + { + "type": "prompt_result", + "id": "old-prompt", + "agentInvoked": True, + "status": "completed", + "sessionSettled": True, + } + ) + case _: + pass + self.respond(frame) + send({"type": "agent_start"}) + match self.scenario: + case scenario if scenario.startswith("chunk-"): + self.invalid_chunk() + case "event-backpressure": + return + case "ui-notify": + send( + { + "type": "extension_ui_request", + "id": "ui-notify", + "method": "notify", + "message": "status", + } + ) + self.end(COMPLETE) + self.result(settled=True) + self.settle() + case "reused-callback": + call = { + "type": "host_tool_call", + "id": "host-dup", + "toolCallId": "call-dup", + "toolName": "band_chat", + "arguments": {}, + } + send(call) + send(call) + case "frame-overflow": + send( + { + "type": "rpc_frame_error", + "originalType": "message_end", + "error": "RPC frame exceeded the transport limit", + } + ) + case "large-image": + message = assistant(COMPLETE) + message["content"] = [ + { + "type": "image", + "mimeType": "image/png", + "data": "YWJj" * 300_000, + } + ] + send( + { + "type": "message_end", + "messageId": "large-image", + "message": message, + } + ) + send({"type": "agent_end", "messages": [message], "isTerminal": True}) + self.result(settled=True) + self.settle() + case "basic": + send({"type": "queue_update", "steering": [], "followUp": []}) + send({"type": "available_commands_update", "commands": []}) + send({"type": "model_changed"}) + send( + { + "type": "tool_execution_start", + "toolCallId": "native-call", + "toolName": "read", + "args": {"path": "note.txt"}, + } + ) + send( + { + "type": "tool_execution_end", + "toolCallId": "native-call", + "toolName": "read", + "result": { + "content": [{"type": "text", "text": "note"}], + "details": {"lines": 1}, + }, + "isError": False, + } + ) + self.end(COMPLETE) + self.result(settled=True) + self.settle() + case "settlement-before-result": + self.end(COMPLETE) + self.settle() + case "no-result": + self.end(COMPLETE) + self.settle() + case scenario if scenario == SETTLED_CALLBACK_SCENARIO: + send( + { + "type": "host_tool_call", + "id": SETTLED_CALLBACK_ID, + "toolCallId": "call-held", + "toolName": "band_chat", + "arguments": {}, + } + ) + self.end(COMPLETE) + self.result(settled=True) + self.settle() + case "host-cancel": + send( + { + "type": "host_tool_call", + "id": HOST_CANCELLED_ID, + "toolCallId": "call-cancelled", + "toolName": "band_chat", + "arguments": {}, + } + ) + case "ui-cancel" | "ui-error": + send( + { + "type": "extension_ui_request", + "id": UI_CANCELLED_ID, + "method": "select", + "title": "Allow command?", + "options": [UI_APPROVE, UI_DENY], + } + ) + case "initial-error" | "initial-aborted": + stop = ( + OmpPromptStatus.ERROR + if self.scenario == "initial-error" + else OmpPromptStatus.ABORTED + ) + self.end("failed", stop=stop) + self.result(settled=True, status=stop) + self.settle() + case _: + self.end(YIELDED, pending=True) + self.result(settled=self.scenario == "false-result-settled") + if self.scenario == "state-settles": + self.settled = True + + def handle(self, frame: Frame) -> None: + global PROTOCOL_VERSION + self.frames.append( + cast(dict[str, JsonValue], frame.model_dump(exclude_unset=True)) + ) + match frame.type: + case OmpCommandName.NEGOTIATE_PROTOCOL: + if self.scenario == "negotiation-rejected": + send( + { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": False, + "error": "Native protocol v2 is not supported", + } + ) + return + self.respond( + frame, + { + "protocolVersion": 1 + if self.scenario == "bad-negotiation" + else OMP_NATIVE_PROTOCOL_VERSION + }, + ) + PROTOCOL_VERSION = OMP_NATIVE_PROTOCOL_VERSION + case OmpCommandName.PROMPT: + self.prompt(frame) + case OmpCommandName.GET_STATE: + self.state_requested = True + self.respond( + frame, + { + "sessionId": "session-local", + "sessionFile": "/local/session.jsonl", + "model": { + "id": "test-model", + "provider": "test-provider", + "name": "Test model", + }, + **( + {"isSettled": self.settled} + if self.scenario != "missing-state" + else {} + ), + "isStreaming": False, + "hasPendingAsyncWork": not self.settled, + "queuedMessageCount": 0, + }, + ) + if ( + self.scenario == SETTLED_CALLBACK_SCENARIO + and self.held is not None + and self.held.type == "checkpoint" + ): + self.checkpoint(self.held) + self.held = None + if self.scenario == "settles-during-state" and not self.settled: + self.settle() + if ( + self.scenario == "event-backpressure" + and self.held is not None + and self.held.type == "event-result" + ): + self.respond(self.held) + self.held = None + case "event-burst": + for number in range(frame.count): + send( + { + "type": "message_update", + "messageId": burst_message_id(number), + "message": {"role": "assistant", "content": str(number)}, + } + ) + self.respond(frame) + case "event-result": + self.held = frame + self.end(COMPLETE, compacted=True) + self.result(settled=True) + self.settle() + case "event-ui": + self.held = frame + send( + { + "type": "extension_ui_request", + "id": "backlog-approval", + "method": "select", + "title": "Allow command?", + "options": [UI_APPROVE, UI_DENY], + } + ) + case "ask-ui": + self.held = frame + send( + { + "type": "extension_ui_request", + "id": "structured-question", + "method": "ask", + "questions": [ + { + "id": "database-choice", + "question": "Which database?", + "options": [{"label": "SQLite"}], + } + ], + } + ) + case "release": + if self.scenario == SETTLED_CALLBACK_SCENARIO: + send( + { + "type": "host_tool_call", + "id": LATE_CALLBACK_ID, + "toolCallId": f"call-{LATE_CALLBACK_ID}", + "toolName": "band_chat", + "arguments": {}, + } + ) + self.respond(frame) + return + if self.scenario == LATE_PROMPT_ERROR: + send( + { + "type": "response", + "id": self.prompt_id, + "command": frame.failed_command or OmpCommandName.PROMPT, + "success": False, + "error": "Prompt command failed after admission", + } + ) + self.respond(frame) + return + if self.scenario == "settlement-before-result": + self.result(settled=False) + else: + send({"type": "agent_start"}) + stop = { + "resume-error": OmpPromptStatus.ERROR, + "resume-error-compacted": OmpPromptStatus.ERROR, + "resume-aborted": OmpPromptStatus.ABORTED, + }.get(self.scenario, "stop") + self.end( + BACKGROUND_COMPLETE, + stop=stop, + compacted=self.scenario == "resume-error-compacted", + ) + self.settle() + self.respond(frame) + case "cancel-host": + send( + { + "type": "host_tool_cancel", + "id": "cancel-frame", + "targetId": HOST_CANCELLED_ID, + } + ) + self.end("host cancelled") + self.result(settled=True) + self.settle() + self.respond(frame) + case "cancel-ui": + send( + { + "type": "extension_ui_request", + "id": "ui-cancel-frame", + "method": "cancel", + "targetId": UI_CANCELLED_ID, + } + ) + self.end("ui cancelled") + self.result(settled=True) + self.settle() + self.respond(frame) + case "host_tool_result" | "extension_ui_response": + if frame.type == "host_tool_result" and frame.id is not None: + self.tool_results.add(frame.id) + if frame.type == "extension_ui_response": + self.ui_answered = True + if ( + self.scenario == "event-backpressure" + and frame.type == "extension_ui_response" + and frame.id == "backlog-approval" + and self.held is not None + and self.held.type == "event-ui" + ): + self.respond(self.held, {"value": frame.value}) + self.held = None + if ( + frame.type == "extension_ui_response" + and frame.id == "structured-question" + and self.held is not None + and self.held.type == "ask-ui" + ): + self.respond( + self.held, + { + "answer": cast( + dict[str, JsonValue], + frame.model_dump(exclude_unset=True), + ) + }, + ) + self.held = None + if ( + self.scenario == "callbacks" + and self.tool_results == {"host-1", "host-2"} + and self.ui_answered + and self.held is not None + ): + self.respond(self.held) + self.held = None + self.end("callbacks complete") + self.result(settled=True) + self.settle() + case "orphan-tool": + send( + { + "type": "host_tool_call", + "id": "host-orphan", + "toolCallId": "call-orphan", + "toolName": "band_chat", + "arguments": {}, + } + ) + case "orphan-ui": + send( + { + "type": "extension_ui_request", + "id": "ui-orphan", + "method": "select", + "title": "Allow command?", + "options": [UI_APPROVE, UI_DENY], + } + ) + self.respond(frame) + case "orphan-ui-cancel": + send( + { + "type": "extension_ui_request", + "id": "ui-cancel-orphan", + "method": "cancel", + } + ) + case "checkpoint": + self.checkpoint(frame) + case OmpCommandName.ABORT: + if self.scenario != "stubborn": + self.respond(frame) + if self.prompt_id is not None: + self.end(OmpPromptStatus.ABORTED, stop=OmpPromptStatus.ABORTED) + self.settle() + case "first" | "hold": + if self.scenario == STDIN_BACKPRESSURE: + self.respond(frame) + threading.Event().wait() + return + self.held = frame + case "second": + self.respond(frame, {"value": "second"}) + if self.held is not None: + self.respond(self.held, {"value": "first"}) + self.held = None + case "error": + send( + { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": False, + "error": "Command failed", + "code": "bad_command", + } + ) + case "crash": + raise SystemExit(CRASH_EXIT) + case "stderr": + sys.stderr.write( + FRAME_PAYLOAD_MARK + * (2 * OMP_FRAME_BYTES // len(FRAME_PAYLOAD_MARK)) + ) + sys.stderr.flush() + self.respond(frame, {"drained": True}) + case _: + self.respond(frame) + + +def announce_ready(scenario: str) -> None: + ready: dict[str, JsonValue] = { + "type": "ready", + "protocolVersion": 1, + "supportedProtocolVersions": [1] + if scenario == "no-v2" + else [1, OMP_NATIVE_PROTOCOL_VERSION], + "maxFrameBytes": OMP_FRAME_BYTES, + "maxReassembledFrameBytes": OMP_REASSEMBLED_BYTES, + } + if scenario == OVERSIZE_AFTER_READY: + ready["maxFrameBytes"] = READY_LIMIT_BYTES + frames = ( + ready, + {"type": "session_settled", "padding": "x" * READY_LIMIT_BYTES}, + ) + # One write keeps the oversized frame buffered behind its ready limits. + sys.stdout.write( + "\n".join(json.dumps(frame, separators=(",", ":")) for frame in frames) + + "\n" + ) + sys.stdout.flush() + return + send(ready) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--scenario", required=True) + parser.add_argument("--gate-port", type=int) + parser.add_argument("--detached-children", action="store_true") + parser.add_argument("--spawn-on-eof", action="store_true") + parser.add_argument( + "--child-stdio", choices=("inherit", "discard"), default="inherit" + ) + args = parser.parse_args() + if args.scenario in {"gated-parent", "gated-child"}: + if args.gate_port is None: + raise ValueError("A process gate port is required") + hold_process_gate( + args.gate_port, + spawn_child=args.scenario == "gated-parent", + detached=args.detached_children, + spawn_on_eof=args.spawn_on_eof, + ) + return + managed_child: subprocess.Popen[bytes] | None = None + if args.scenario in { + "descendants", + "no-ready-descendants", + "lifecycle", + "eof-descendants", + }: + if args.gate_port is None: + raise ValueError("A process gate port is required") + # The gate, not the peer's stdio EOF, owns these children's lifetime. + managed_child = start_descendant( + "gated-parent", + args.gate_port, + args.child_stdio, + detached=args.detached_children, + spawn_on_eof=args.scenario == "eof-descendants", + ) + if args.scenario == "stubborn" and hasattr(signal, "SIGTERM"): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + peer = Peer(args.scenario) + if args.scenario not in {"no-ready", "no-ready-descendants"}: + announce_ready(args.scenario) + for line in sys.stdin: + peer.handle(Frame.model_validate_json(line, strict=True)) + if managed_child is not None and managed_child.stdin is not None: + managed_child.stdin.close() + if args.scenario == "stubborn": + threading.Event().wait() + + +if __name__ == "__main__": + main() diff --git a/tests/integrations/ompprocess/__init__.py b/tests/integrations/ompprocess/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py new file mode 100644 index 000000000..629a5c1ba --- /dev/null +++ b/tests/integrations/ompprocess/test_launch.py @@ -0,0 +1,24 @@ +"""Native OMP launch settings: command shape and provider environment.""" + +from __future__ import annotations + +from band.integrations.ompprocess import omp_runtime_env + + +class TestNativeLaunchSettings: + def test_runtime_env_routes_the_selected_provider_key(self) -> None: + assert omp_runtime_env(env=None, model="openai/custom", api_key="secret") == { + "OPENAI_API_KEY": "secret" + } + + def test_explicit_env_overrides_the_derived_credential(self) -> None: + env = omp_runtime_env( + env={"OPENAI_API_KEY": "override", "EXTRA": "1"}, + model="openai/custom", + api_key="secret", + ) + assert env == {"OPENAI_API_KEY": "override", "EXTRA": "1"} + + def test_no_credentials_and_no_env_means_inherit(self) -> None: + assert omp_runtime_env(env=None, model=None, api_key=None) is None + assert omp_runtime_env(env={}, model="openai/custom", api_key=None) is None diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py new file mode 100644 index 000000000..f800d03af --- /dev/null +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -0,0 +1,204 @@ +"""The SDK-owned plugin projects MCP servers and connection rules for one process.""" + +from __future__ import annotations + +import json +import os +import stat +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from band.integrations.ompprocess import ( + OMP_MCP_SERVERS_CONFIG, + OMP_MCP_STARTUP_CONFIG, + OMP_MCP_STARTUP_JS, + RELATIVE_MCP_COMMAND_PREFIXES, + OmpMcpServer, + OmpMcpTransport, + omp_mcp_plugin, +) +from tests.ostable import skip_unless_os + + +def plugin_privacy(root: Path, files: tuple[Path, ...]) -> tuple[int, int | None]: + file_modes = {stat.S_IMODE(path.stat().st_mode) for path in files} + file_mode = next(iter(file_modes)) if len(file_modes) == 1 else None + return (stat.S_IMODE(root.stat().st_mode), file_mode) + + +class TestMcpPluginProjection: + @pytest.mark.parametrize("system_prompt", ["", " "]) + def test_nothing_to_project_yields_no_plugin( + self, tmp_path: Path, system_prompt: str + ) -> None: + with omp_mcp_plugin( + {}, workspace=str(tmp_path), system_prompt=system_prompt + ) as plugin: + assert plugin is None + + def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: + servers = { + "local": OmpMcpServer(command="./server", args=("--port", "1")), + "remote": OmpMcpServer( + type=OmpMcpTransport.SSE, url="https://mcp.example/sse" + ), + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + + config = json.loads( + (root / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + startup_config = json.loads( + (root / OMP_MCP_STARTUP_CONFIG).read_text(encoding="utf-8") + ) + assert startup_config == {"startup": {"quiet": False}} + assert config["mcpServers"]["remote"] == { + "type": OmpMcpTransport.SSE, + "url": "https://mcp.example/sse", + } + local = config["mcpServers"]["local"] + assert local["args"] == ["--port", "1"] + assert Path(local["cwd"]) == tmp_path.resolve() + assert Path(local["command"]) == (tmp_path / "server").resolve() + + assert not Path(plugin).exists() + + def test_connection_instructions_become_an_always_applied_rule( + self, tmp_path: Path + ) -> None: + with omp_mcp_plugin( + {}, workspace=str(tmp_path), system_prompt="Reply through Band." + ) as plugin: + assert plugin is not None + rules = list((Path(plugin) / "rules").glob("band-connection-*.md")) + assert len(rules) == 1 + assert rules[0].read_text(encoding="utf-8") == ( + "---\nalwaysApply: true\n---\nReply through Band." + ) + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + def test_projected_plugin_files_are_owner_only( + self, tmp_path: Path, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + servers = {"local": OmpMcpServer(command="./server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + files = ( + root / OMP_MCP_SERVERS_CONFIG, + root / OMP_MCP_STARTUP_CONFIG, + root / OMP_MCP_STARTUP_JS, + ) + assert root.is_dir() + assert all(path.is_file() for path in files) + assert plugin_privacy(root, files) == (0o700, 0o600) + + @pytest.mark.parametrize("accepted_on", [("nt",)]) + def test_projected_plugin_files_are_readable( + self, tmp_path: Path, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + servers = {"local": OmpMcpServer(command="./server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + files = ( + root / OMP_MCP_SERVERS_CONFIG, + root / OMP_MCP_STARTUP_CONFIG, + root / OMP_MCP_STARTUP_JS, + ) + assert root.is_dir() + assert all(path.is_file() for path in files) + assert all(os.access(path, os.R_OK) for path in files) + + @pytest.mark.parametrize("prefix", RELATIVE_MCP_COMMAND_PREFIXES) + def test_relative_command_is_absolutized(self, tmp_path: Path, prefix: str) -> None: + servers = {"local": OmpMcpServer(command=f"{prefix}server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + assert Path(config["mcpServers"]["local"]["command"]).is_absolute() + + @pytest.mark.parametrize("transport", [None, OmpMcpTransport.STDIO]) + def test_command_transport_projects_a_launchable_server( + self, tmp_path: Path, transport: OmpMcpTransport | None + ) -> None: + servers = { + "local": OmpMcpServer( + type=transport, command="server", args=("--port", "1") + ) + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + local = config["mcpServers"]["local"] + assert local.get("type") == transport + assert local["command"] == "server" + assert local["args"] == ["--port", "1"] + assert Path(local["cwd"]) == tmp_path.resolve() + + @pytest.mark.parametrize( + "transport", [None, OmpMcpTransport.HTTP, OmpMcpTransport.SSE] + ) + def test_url_transport_preserves_the_remote_endpoint( + self, tmp_path: Path, transport: OmpMcpTransport | None + ) -> None: + servers = { + "remote": OmpMcpServer( + type=transport, + url="https://mcp.example/tools", + headers={"Authorization": "Bearer token"}, + ) + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + remote = config["mcpServers"]["remote"] + assert remote.get("type") == transport + assert remote["url"] == "https://mcp.example/tools" + assert remote["headers"] == {"Authorization": "Bearer token"} + assert "command" not in remote + assert "cwd" not in remote + + @pytest.mark.parametrize( + "fields", + [ + pytest.param({}, id="no-endpoint"), + pytest.param( + {"command": "server", "url": "https://mcp.example"}, + id="two-endpoints", + ), + pytest.param( + {"command": "server", "type": OmpMcpTransport.HTTP}, + id="command-http", + ), + pytest.param( + {"command": "server", "type": OmpMcpTransport.SSE}, + id="command-sse", + ), + pytest.param( + {"url": "https://mcp.example", "type": OmpMcpTransport.STDIO}, + id="url-stdio", + ), + pytest.param( + {"command": "server", "type": "websocket"}, + id="unknown-transport", + ), + ], + ) + def test_conflicting_transport_settings_are_rejected( + self, fields: dict[str, str] + ) -> None: + with pytest.raises(ValidationError): + OmpMcpServer.model_validate(fields) diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py new file mode 100644 index 000000000..df5f6ae1e --- /dev/null +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -0,0 +1,180 @@ +"""Required MCP connections cannot be inferred from an otherwise ready OMP process.""" + +from __future__ import annotations + +import json +from collections.abc import Awaitable, Callable +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +pytest.importorskip("psutil", reason="psutil not installed (band-sdk[omp])") + +from band.integrations.ompprocess import ( + OMP_MCP_RECEIPT, + OMP_PLUGIN_ROOT, + McpConnectionStatus, + OmpMcpServer, + OmpStartupError, + omp_mcp_ready, + omp_mcp_ready_command, +) +from band.integrations.omprpc import OmpRpcClient +from band.integrations.ompwire import ( + OmpEvent, + OmpEventType, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +REQUIRED_SERVER = f"{OMP_PLUGIN_ROOT}:required" + + +class TestRequiredMcpStartup: + @pytest.mark.parametrize( + "receipt", + [ + pytest.param(None, id="missing-receipt"), + pytest.param("not JSON", id="malformed-receipt"), + pytest.param( + {"systemPromptReady": True, "servers": []}, + id="missing-server", + ), + pytest.param( + { + "systemPromptReady": True, + "servers": [ + { + "type": McpConnectionStatus.CONNECTED, + "serverName": REQUIRED_SERVER, + }, + { + "type": McpConnectionStatus.CONNECTED, + "serverName": REQUIRED_SERVER, + }, + ], + }, + id="duplicate-server", + ), + pytest.param( + { + "systemPromptReady": True, + "servers": [ + { + "type": McpConnectionStatus.CONNECTED, + "serverName": "unrelated", + } + ], + }, + id="unexpected-server", + ), + pytest.param( + { + "systemPromptReady": True, + "servers": [ + { + "type": McpConnectionStatus.FAILED, + "serverName": REQUIRED_SERVER, + } + ], + }, + id="missing-failure-detail", + ), + pytest.param( + { + "systemPromptReady": True, + "servers": [ + { + "type": McpConnectionStatus.FAILED, + "serverName": REQUIRED_SERVER, + "error": "initialize rejected", + } + ], + }, + id="failed-server", + ), + pytest.param( + { + "systemPromptReady": False, + "servers": [ + { + "type": McpConnectionStatus.CONNECTED, + "serverName": REQUIRED_SERVER, + } + ], + }, + id="unloaded-instructions", + ), + ], + ) + async def test_rejects_incomplete_or_failed_receipts( + self, + tmp_path: Path, + receipt: str | dict[str, object] | None, + ) -> None: + plugin = tmp_path / OMP_PLUGIN_ROOT + plugin.mkdir() + if receipt is not None: + payload = json.dumps(receipt) if isinstance(receipt, dict) else receipt + (plugin / OMP_MCP_RECEIPT).write_text(payload, encoding="utf-8") + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = { + "commands": [{"name": omp_mcp_ready_command(str(plugin))}], + } + with pytest.raises(OmpStartupError): + await omp_mcp_ready( + runtime, + str(plugin), + {"required": OmpMcpServer(command="server")}, + ) + + async def test_missing_startup_extension_is_rejected(self, tmp_path: Path) -> None: + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = {"commands": []} + with pytest.raises(OmpStartupError): + await omp_mcp_ready(runtime, str(tmp_path), {}) + + @pytest.mark.parametrize("invoke", ["event", "tool", "ui"]) + async def test_startup_callbacks_cannot_fall_through_to_inference( + self, tmp_path: Path, invoke: str + ) -> None: + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = { + "commands": [{"name": omp_mcp_ready_command(str(tmp_path))}], + } + + async def prompt( + message: str, + *, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + _ = message + match invoke: + case "event": + await on_event( + OmpEvent.model_validate({"type": OmpEventType.AGENT_START}) + ) + case "tool": + await on_tool( + OmpHostToolCall.model_validate( + { + "id": "host-1", + "toolCallId": "call-1", + "toolName": "band_chat", + "arguments": {}, + } + ) + ) + case "ui": + await on_ui( + OmpUIRequest.model_validate({"id": "ask", "method": "select"}) + ) + + runtime.prompt.side_effect = prompt + with pytest.raises(OmpStartupError): + await omp_mcp_ready(runtime, str(tmp_path), {}) diff --git a/tests/integrations/ompprocess/test_tool_names.py b/tests/integrations/ompprocess/test_tool_names.py new file mode 100644 index 000000000..ddddc7f5e --- /dev/null +++ b/tests/integrations/ompprocess/test_tool_names.py @@ -0,0 +1,59 @@ +"""Tool names OMP mints under the SDK plugin namespace map back to Band tools.""" + +from __future__ import annotations + +import pytest + +from band.integrations.omp import XD_MCP_PREFIX, XD_URL_PREFIX +from band.integrations.ompprocess import ( + OMP_PLUGIN_TOOL_NAMESPACE, + normalize_native_mcp_device_call, + normalize_native_mcp_tool_name, +) + +OWN_TOOL = "band_send_message" +OWN = frozenset({OWN_TOOL}) +BARE_TITLE = f"{XD_MCP_PREFIX}{OWN_TOOL}" +NAMESPACED_TITLE = f"{XD_MCP_PREFIX}{OMP_PLUGIN_TOOL_NAMESPACE}band_{OWN_TOOL}" +BARE_PATH = f"{XD_URL_PREFIX}{BARE_TITLE}" +NAMESPACED_PATH = f"{XD_URL_PREFIX}{NAMESPACED_TITLE}" + + +class TestPluginNamespacedToolNames: + @pytest.mark.parametrize( + "path", + [ + BARE_PATH, + NAMESPACED_PATH, + ], + ) + def test_executed_device_write_maps_to_the_band_tool(self, path: str) -> None: + arguments = {"path": path, "content": '{"chat_id":"room-1","content":"hi"}'} + assert normalize_native_mcp_device_call( + "write", arguments, OWN, kind="execute" + ) == (OWN_TOOL, {"chat_id": "room-1", "content": "hi"}) + + @pytest.mark.parametrize("title", [BARE_TITLE, NAMESPACED_TITLE]) + def test_mcp_title_maps_to_the_band_tool(self, title: str) -> None: + assert normalize_native_mcp_tool_name(title, OWN) == OWN_TOOL + assert normalize_native_mcp_device_call(title, {"chat_id": "r1"}, OWN) == ( + OWN_TOOL, + {"chat_id": "r1"}, + ) + + def test_a_namespaced_read_is_never_the_band_tool(self) -> None: + arguments = { + "path": NAMESPACED_PATH, + "content": '{"content":"hi"}', + } + assert normalize_native_mcp_device_call( + "read", arguments, OWN, kind="read" + ) == ("read", arguments) + + def test_a_foreign_namespace_stays_unmapped(self) -> None: + assert ( + normalize_native_mcp_tool_name( + f"{XD_MCP_PREFIX}other_sdk_band_send_message", OWN + ) + is None + ) diff --git a/tests/integrations/omprpc/__init__.py b/tests/integrations/omprpc/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integrations/omprpc/backpressure.py b/tests/integrations/omprpc/backpressure.py new file mode 100644 index 000000000..1597758a0 --- /dev/null +++ b/tests/integrations/omprpc/backpressure.py @@ -0,0 +1,35 @@ +"""Observe real subprocess pipe saturation without replacing writes or drains.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from band.integrations.ompwire import OMP_FRAME_BYTES + +PRESSURE_WRITE_COUNT = 8 + + +class BackpressureProtocol(asyncio.subprocess.SubprocessStreamProtocol): + def __init__( + self, limit: int, loop: asyncio.AbstractEventLoop, paused: asyncio.Event + ) -> None: + super().__init__(limit, loop) + self.paused = paused + + def pause_writing(self) -> None: + super().pause_writing() + self.paused.set() + + +class BackpressureSpawn: + def __init__(self) -> None: + self.paused = asyncio.Event() + + async def create( + self, *command: str, limit: int = OMP_FRAME_BYTES, **kwargs: Any + ) -> asyncio.subprocess.Process: + loop = asyncio.get_running_loop() + protocol = BackpressureProtocol(limit, loop, self.paused) + transport, _ = await loop.subprocess_exec(lambda: protocol, *command, **kwargs) + return asyncio.subprocess.Process(transport, protocol, loop) diff --git a/tests/integrations/omprpc/cancellation.py b/tests/integrations/omprpc/cancellation.py new file mode 100644 index 000000000..82621a3eb --- /dev/null +++ b/tests/integrations/omprpc/cancellation.py @@ -0,0 +1,48 @@ +"""Hold real RPC cleanup and observe cancellation-safe joins without elapsed time.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable + +from band.integrations.omprpc import OmpRpcClient + + +class HeldCleanup: + def __init__(self, client: OmpRpcClient) -> None: + self.client = client + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.finished = asyncio.Event() + self.task: asyncio.Task[object] | None = None + self.joiners: asyncio.Queue[asyncio.Task[object] | None] = asyncio.Queue() + closer = client.close_process + join_close = client.join_close + + async def hold_cleanup() -> None: + self.task = asyncio.current_task() + self.entered.set() + await self.release.wait() + await closer() + self.finished.set() + + def observe_join(operation: Awaitable[None]) -> Awaitable[None]: + waiter = join_close(operation) + self.joiners.put_nowait(asyncio.current_task()) + return waiter + + client.close_process = hold_cleanup + client.join_close = observe_join + + async def wait_for_join(self, caller: asyncio.Task[None]) -> None: + async with asyncio.timeout(5): + await self.entered.wait() + joining = await self.joiners.get() + assert joining is caller + assert self.task is not None + assert not self.task.done() + assert not caller.done() + + async def cancel_and_wait_for_rejoin(self, caller: asyncio.Task[None]) -> None: + assert caller.cancel() + await self.wait_for_join(caller) diff --git a/tests/integrations/omprpc/conftest.py b/tests/integrations/omprpc/conftest.py new file mode 100644 index 000000000..a24840303 --- /dev/null +++ b/tests/integrations/omprpc/conftest.py @@ -0,0 +1,50 @@ +"""Client lifecycle and process gate fixtures for OMP RPC tests.""" + +from __future__ import annotations + +import sys +from collections.abc import AsyncIterator + +import pytest +import pytest_asyncio + +pytest.importorskip("psutil", reason="psutil not installed (band-sdk[omp])") + +from band.integrations.omprpc import OmpRpcClient +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ClientFactory, + ProcessGate, + process_gate_session, +) + + +@pytest_asyncio.fixture(loop_scope="function") +async def client_factory() -> AsyncIterator[ClientFactory]: + clients: list[OmpRpcClient] = [] + + async def build(scenario: str) -> OmpRpcClient: + client = OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + scenario, + ) + ) + clients.append(client) + await client.start() + return client + + try: + yield build + finally: + for client in clients: + await client.close() + + +@pytest_asyncio.fixture(loop_scope="function") +async def process_gate() -> AsyncIterator[ProcessGate]: + async with process_gate_session() as gate: + yield gate diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py new file mode 100644 index 000000000..6512ddf62 --- /dev/null +++ b/tests/integrations/omprpc/kit.py @@ -0,0 +1,551 @@ +"""Shared observers and process gates for OMP RPC tests.""" + +from __future__ import annotations + +import asyncio +import os +import signal +import sys +import threading +import time +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass +from typing import Any, cast + +import psutil +import pytest +from pydantic import JsonValue + +from band.integrations.codex.winspawn import Job +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.ompwire import ( + OmpCommandName, + OmpContent, + OmpEvent, + OmpEventType, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations import omppeer +from tests.integrations.omppeer import COMPLETE, FOLLOW_UP_PROMPT, ProcessReceipt +from tests.integrations.omprpc.cancellation import HeldCleanup + +ClientFactory = Callable[[str], Awaitable[OmpRpcClient]] + +# The fake OMP peer every RPC test launches as its subprocess. +OMP_PEER_SCRIPT = omppeer.__file__ + + +class FailingOwnedClient(OmpRpcClient): + fail_spawn = True + + async def spawn_owned_process( + self, *command: str, **kwargs: Any + ) -> asyncio.subprocess.Process: + process = await super().spawn_owned_process(*command, **kwargs) + if self.fail_spawn: + raise ValueError("Owned spawn wrapper failed") + return process + + +class SpawnBarrier: + """Create the first subprocess, then hold returning it until released.""" + + def __init__( + self, spawn: Callable[..., Awaitable[asyncio.subprocess.Process]] + ) -> None: + self.spawn = spawn + self.created = asyncio.Event() + self.release = asyncio.Event() + self.process: asyncio.subprocess.Process | None = None + + async def create( + self, *command: str, **kwargs: object + ) -> asyncio.subprocess.Process: + process = await self.spawn(*command, **kwargs) + if self.process is None: + self.process = process + self.created.set() + await self.release.wait() + return process + + +class Observer: + def __init__(self) -> None: + self.events: list[OmpEvent] = [] + self.queue: asyncio.Queue[OmpEvent] = asyncio.Queue() + + @property + def assistant_texts(self) -> list[str]: + return [ + part.text + for event in self.events + if event.message is not None + and event.message.role == "assistant" + and isinstance(event.message.content, list) + for part in event.message.content + if part.text is not None + ] + + @property + def outline(self) -> list[OmpEventType]: + return [event.type for event in self.events] + + @property + def update_ids(self) -> list[str | None]: + return [ + event.messageId + for event in self.events + if event.type == OmpEventType.MESSAGE_UPDATE + ] + + async def on_event(self, event: OmpEvent) -> None: + self.events.append(event) + self.queue.put_nowait(event) + + async def on_tool(self, call: OmpHostToolCall) -> OmpToolResult: + return OmpToolResult(content=[OmpContent(type="text", text=call.toolName)]) + + async def on_ui(self, request: OmpUIRequest) -> OmpUIResponse: + return OmpUIResponse(cancelled=True) + + async def wait_for(self, event_type: OmpEventType) -> OmpEvent: + async with asyncio.timeout(5): + while True: + event = await self.queue.get() + if event.type == event_type: + return event + + +def prompt_task( + client: OmpRpcClient, observer: Observer, message: str = "first" +) -> asyncio.Task[None]: + return asyncio.create_task( + client.prompt( + message, + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + ) + + +def owned_rpc_process(client: OmpRpcClient) -> asyncio.subprocess.Process: + process = client._proc + assert process is not None + return process + + +async def second_prompt(client: OmpRpcClient) -> Observer: + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer, FOLLOW_UP_PROMPT) + assert COMPLETE in observer.assistant_texts + return observer + + +async def announced_prompt( + client: OmpRpcClient, + observer: Observer, + entered: asyncio.Event, + message: str, +) -> None: + entered.set() + await client.prompt( + message, + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + + +async def queue_prompt( + client: OmpRpcClient, message: str = "queued" +) -> asyncio.Task[None]: + entered = asyncio.Event() + task = asyncio.create_task(announced_prompt(client, Observer(), entered, message)) + # The coroutine reaches the occupied admission lock before this waiter resumes. + await entered.wait() + return task + + +def descendant_client( + gate: ProcessGate, + *, + scenario: str = "descendants", + stdio: str = "discard", + detached: bool = False, +) -> OmpRpcClient: + return OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + scenario, + "--gate-port", + str(gate.port), + "--child-stdio", + stdio, + *(("--detached-children",) if detached else ()), + ) + ) + + +async def checkpoint_frames(client: OmpRpcClient) -> list[dict[str, JsonValue]]: + checkpoint = await client.request("checkpoint") + return cast(list[dict[str, JsonValue]], checkpoint["frames"]) + + +async def sent_prompt_messages(client: OmpRpcClient) -> list[str]: + return [ + cast(str, frame["message"]) + for frame in await checkpoint_frames(client) + if frame["type"] == OmpCommandName.PROMPT + ] + + +def process_is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def is_descendant(child: psutil.Process, ancestor: psutil.Process) -> bool: + """True when ``ancestor`` is a living parent of ``child``. + + Windows Python launchers insert an extra process, so the gate child's + ``parent_pid`` is not always the owned helper. + """ + try: + ancestor_pid = ancestor.pid + ancestor_created = ancestor.create_time() + current = child.parent() + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + seen: set[int] = set() + while current is not None: + try: + pid = current.pid + if pid == ancestor_pid and current.create_time() == ancestor_created: + return True + if pid in seen: + return False + seen.add(pid) + current = current.parent() + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + return False + + +def kill_gate_processes(processes: list[psutil.Process]) -> None: + for process in processes: + with suppress(psutil.NoSuchProcess): + process.kill() + deadline = time.monotonic() + 5 + while alive := [process.pid for process in processes if process_is_alive(process)]: + if time.monotonic() >= deadline: + raise RuntimeError(f"Process gate cleanup left live processes: {alive}") + time.sleep(0.01) + + +@dataclass(frozen=True) +class GatedProcess: + receipt: ProcessReceipt + process: psutil.Process + + +async def descendant_liveness(descendants: Sequence[GatedProcess]) -> list[bool]: + return [ + await asyncio.to_thread(process_is_alive, child.process) + for child in descendants + ] + + +class ProcessGate: + def __init__(self) -> None: + self.port = 0 + self.connections: asyncio.Queue[ + tuple[asyncio.StreamReader, asyncio.StreamWriter] + ] = asyncio.Queue() + self.writers: list[asyncio.StreamWriter] = [] + self.processes: list[psutil.Process] = [] + + def accept( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + self.writers.append(writer) + self.connections.put_nowait((reader, writer)) + + async def connection(self) -> GatedProcess: + reader, _ = await self.connections.get() + receipt = ProcessReceipt.model_validate_json(await reader.readline()) + process = await asyncio.to_thread(psutil.Process, receipt.pid) + self.processes.append(process) + assert await asyncio.to_thread(process.create_time) == receipt.created + return GatedProcess(receipt, process) + + +@asynccontextmanager +async def process_gate_session() -> AsyncIterator[ProcessGate]: + gate = ProcessGate() + server = await asyncio.start_server(gate.accept, "127.0.0.1", 0) + try: + sockets = server.sockets + assert sockets + gate.port = sockets[0].getsockname()[1] + yield gate + finally: + server.close() + for writer in gate.writers: + writer.close() + await asyncio.gather( + *(writer.wait_closed() for writer in gate.writers), return_exceptions=True + ) + await server.wait_closed() + await asyncio.to_thread(kill_gate_processes, gate.processes) + + +def refuse_process_signal(*args: object) -> None: + raise PermissionError("Process ownership is temporarily not signalable") + + +def refuse_containment(monkeypatch: pytest.MonkeyPatch) -> None: + if os.name == "nt": + monkeypatch.setattr(Job, "terminate", refuse_process_signal) + return + monkeypatch.setattr(os, "killpg", refuse_process_signal) + + +def session_ids(pids: Sequence[int]) -> list[int]: + getsid = getattr(os, "getsid", None) + assert getsid is not None + return [getsid(pid) for pid in pids] + + +def job_member_pids(client: OmpRpcClient) -> set[int]: + owner = client._windows_process + assert owner is not None and owner.job is not None + return set(owner.job._process_ids()) + + +async def wait_for_descendants(gate: ProcessGate, count: int = 2) -> list[GatedProcess]: + return [await gate.connection() for _ in range(count)] + + +async def descendants_started_with( + gate: ProcessGate, root: asyncio.subprocess.Process, count: int = 2 +) -> list[GatedProcess]: + ancestor = await asyncio.to_thread(psutil.Process, root.pid) + children: list[GatedProcess] = [] + while len(children) < count: + child = await gate.connection() + if is_descendant(child.process, ancestor): + children.append(child) + return children + + +async def crash_rpc_root(client: OmpRpcClient) -> None: + root = owned_rpc_process(client) + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await root.wait() + assert root.returncode == omppeer.CRASH_EXIT + with pytest.raises(OmpProcessError): + await crashing + + +@dataclass +class DescendantTree: + client: OmpRpcClient + cleanup: HeldCleanup + starting: asyncio.Task[None] + root: asyncio.subprocess.Process + descendants: list[GatedProcess] + + +@asynccontextmanager +async def hold_after_stdin_closed( + client: OmpRpcClient, +) -> AsyncIterator[asyncio.Task[None]]: + entered = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def hold() -> None: + loop.call_soon_threadsafe(entered.set) + if not release.wait(timeout=10): + raise TimeoutError("Test did not release after_stdin_closed") + + client.after_stdin_closed = hold + closing = asyncio.create_task(client.close()) + try: + async with asyncio.timeout(10): + await entered.wait() + yield closing + if not closing.done(): + release.set() + await closing + finally: + release.set() + await asyncio.gather(closing, return_exceptions=True) + + +async def start_eof_client( + gate: ProcessGate, *, stdio: str +) -> tuple[OmpRpcClient, GatedProcess, asyncio.subprocess.Process]: + client = descendant_client( + gate, scenario="eof-descendants", stdio=stdio, detached=True + ) + await client.start() + helper = await gate.connection() + return client, helper, owned_rpc_process(client) + + +async def connect_descendant_tree( + gate: ProcessGate, *, scenario: str, stdio: str +) -> DescendantTree: + client = descendant_client(gate, scenario=scenario, stdio=stdio) + cleanup = HeldCleanup(client) + starting = asyncio.create_task(client.start()) + descendants = await wait_for_descendants(gate) + return DescendantTree( + client=client, + cleanup=cleanup, + starting=starting, + root=owned_rpc_process(client), + descendants=descendants, + ) + + +class DeferredSessionSignal: + """Keep real members alive while post-signal ownership lookup disappears.""" + + def __init__( + self, + pgid: int, + members: Sequence[GatedProcess], + *, + hide_after_signal: bool, + ) -> None: + self.pgid = pgid + self.members = {member.process.pid for member in members} + self.hide_after_signal = hide_after_signal + self.signalled = False + self.original_group = os.getpgid + self.original_kill = os.kill + self.original_environment = psutil.Process.environ + + def send(self, pgid: int, sig: int) -> None: + assert pgid == self.pgid + assert sig == signal.SIGKILL + self.signalled = True + + def send_pid(self, pid: int, sig: int) -> None: + if sig == signal.SIGKILL: + assert pid in self.members + self.signalled = True + return + self.original_kill(pid, sig) + + def group(self, pid: int) -> int: + if self.signalled and self.hide_after_signal and pid in self.members: + raise ProcessLookupError(pid) + return self.original_group(pid) + + def environment(self, process: psutil.Process) -> dict[str, str]: + if self.signalled and self.hide_after_signal and process.pid in self.members: + raise psutil.NoSuchProcess(process.pid) + return self.original_environment(process) + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr(os, "killpg", self.send) + patch.setattr(os, "kill", self.send_pid) + patch.setattr(os, "getpgid", self.group) + patch.setattr( + psutil.Process, "environ", lambda process: self.environment(process) + ) + + +@asynccontextmanager +async def unrelated_gate_child(gate: ProcessGate) -> AsyncIterator[GatedProcess]: + process = await asyncio.create_subprocess_exec( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + "gated-child", + "--gate-port", + str(gate.port), + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + start_new_session=os.name != "nt", + ) + try: + yield await gate.connection() + finally: + if process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + async with asyncio.timeout(5): + await process.wait() + + +class RecycledGroup: + """Reuse a reaped leader's group id during membership enumeration.""" + + def __init__( + self, root: asyncio.subprocess.Process, unrelated: GatedProcess + ) -> None: + self.root = root + self.unrelated = unrelated + self.reused = False + + def check_leader(self, pid: int, sig: int) -> None: + assert (pid, sig) == (self.root.pid, 0) + if not self.reused: + raise ProcessLookupError + + def group(self, pid: int) -> int: + assert pid == self.unrelated.process.pid + self.reused = True + return self.root.pid + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr(psutil, "pids", lambda: [self.unrelated.process.pid]) + patch.setattr(os, "kill", self.check_leader) + patch.setattr(os, "getpgid", self.group) + + +class RecycledEnvironment: + """Expose one stale matching environment while the process identity changes.""" + + def __init__(self, unrelated: GatedProcess, owner_pid: int) -> None: + self.pid = unrelated.process.pid + self.owner_environment = psutil.Process(owner_pid).environ() + self.original_environment = psutil.Process.environ + self.original_create_time = psutil.Process.create_time + self.recycled = False + + def environment(self, process: psutil.Process) -> dict[str, str]: + if process.pid == self.pid and not self.recycled: + self.recycled = True + return self.owner_environment + return self.original_environment(process) + + def create_time(self, process: psutil.Process) -> float: + created = self.original_create_time(process) + return created + 1 if self.recycled and process.pid == self.pid else created + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr( + psutil.Process, "environ", lambda process: self.environment(process) + ) + patch.setattr( + psutil.Process, "create_time", lambda process: self.create_time(process) + ) diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py new file mode 100644 index 000000000..1f2fb0cad --- /dev/null +++ b/tests/integrations/omprpc/lifecycle.py @@ -0,0 +1,290 @@ +"""A live ``OmpRpcClient`` driven step by step against a real RPC peer. + +This is the model for ``test_lifecycle_model.py``, not a collected test. +Hypothesis picks a short sequence of ``Step`` values — the operations a host +can issue — and ``RpcLifecycle.take`` applies each one to one client and the +process tree it owns. After the sequence, the client must still start, finish +a turn, and reap descendants. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Awaitable +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass +from enum import StrEnum +from unittest.mock import patch + +from hypothesis import strategies as st + +from band.integrations.omprpc import OmpProcessError +from band.integrations.ompwire import OmpEventType +from tests.integrations.omppeer import BACKGROUND_COMPLETE, CRASH_EXIT +from tests.integrations.omprpc.kit import ( + GatedProcess, + Observer, + ProcessGate, + SpawnBarrier, + descendant_client, + descendant_liveness, + descendants_started_with, + owned_rpc_process, + process_gate_session, + prompt_task, + queue_prompt, + second_prompt, +) + +# A healthy step takes milliseconds; a hung one fails its example within this. +STEP_GUARD_S = 15 + + +class Step(StrEnum): + """Host operations Hypothesis may interleave.""" + + PROMPT = "admit a held prompt" + CANCEL_QUEUED = "cancel a prompt waiting for admission" + SETTLE = "settle the admitted prompt" + CANCEL_OWNER = "cancel the admitted prompt" + CLOSE = "close the transport" + CRASH = "crash the root before closing" + CANCEL_START_RESTART = "cancel held spawn while a restart waits" + CANCEL_CLOSE_RESTART = "cancel close during held spawn while a restart waits" + + +sequences = st.lists(st.sampled_from(Step), min_size=1, max_size=6) + + +@dataclass +class HeldSpawn: + starting: asyncio.Task[None] + root: asyncio.subprocess.Process + children: list[GatedProcess] + barrier: SpawnBarrier + + def release(self) -> None: + self.barrier.release.set() + + +async def announced_operation( + operation: Awaitable[None], entered: asyncio.Event +) -> None: + entered.set() + await operation + + +class RpcLifecycle: + """One RPC client, the process trees it has owned, and at most one admitted turn.""" + + def __init__(self, gate: ProcessGate, *, detached: bool) -> None: + self.gate = gate + self.client = descendant_client(gate, scenario="lifecycle", detached=detached) + self.live = False + self.active: asyncio.Task[None] | None = None + self.observer: Observer | None = None + self.tasks: list[asyncio.Task[None]] = [] + self.roots: list[asyncio.subprocess.Process] = [] + self.children: list[GatedProcess] = [] + + async def record_children(self) -> list[GatedProcess]: + children = await descendants_started_with(self.gate, self.roots[-1]) + self.children.extend(children) + return children + + async def open(self) -> None: + if self.live: + return + await self.client.start() + self.roots.append(owned_rpc_process(self.client)) + await self.record_children() + self.live = True + + async def hold_prompt(self) -> None: + await self.open() + if self.active is not None: + return + self.observer = Observer() + self.active = prompt_task(self.client, self.observer) + self.tasks.append(self.active) + await self.observer.wait_for(OmpEventType.AGENT_END) + assert not self.active.done() + assert (await self.client.request("get_state"))["isSettled"] is False + + async def cancel_queued(self) -> None: + await self.hold_prompt() + queued = await queue_prompt(self.client) + self.tasks.append(queued) + queued.cancel() + [result] = await asyncio.gather(queued, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + assert self.active is not None and not self.active.done() + assert (await self.client.request("get_state"))["isSettled"] is False + + async def settle(self) -> None: + await self.hold_prompt() + await self.client.request("release") + assert self.active is not None and self.observer is not None + await self.active + assert BACKGROUND_COMPLETE in self.observer.assistant_texts + self.active = None + assert (await self.client.request("get_state"))["isSettled"] is True + + async def cancel_owner(self) -> None: + await self.hold_prompt() + assert self.active is not None + self.active.cancel() + [result] = await asyncio.gather(self.active, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + self.active = None + await self.close() + + async def crash(self) -> None: + await self.open() + root = owned_rpc_process(self.client) + with suppress(OmpProcessError): + await self.client.request("crash") + await root.wait() + assert root.returncode == CRASH_EXIT + await self.close() + + async def assert_reaped( + self, roots: list[asyncio.subprocess.Process], children: list[GatedProcess] + ) -> None: + assert all(root.returncode is not None for root in roots) + assert not any(await descendant_liveness(children)) + + async def close(self) -> None: + await self.client.close() + if self.active is not None: + [result] = await asyncio.gather(self.active, return_exceptions=True) + assert result is None or isinstance(result, OmpProcessError) + self.active = None + self.live = False + await self.assert_reaped(self.roots, self.children) + + def _track(self, task: asyncio.Task[None]) -> asyncio.Task[None]: + self.tasks.append(task) + return task + + async def _assert_finished( + self, task: asyncio.Task[None], error: type[BaseException] + ) -> None: + [result] = await asyncio.gather(task, return_exceptions=True) + assert isinstance(result, error) + + async def _queue_close(self) -> asyncio.Task[None]: + entered = asyncio.Event() + closing = self._track( + asyncio.create_task(announced_operation(self.client.close(), entered)) + ) + await entered.wait() + return closing + + async def _queue_restart(self) -> asyncio.Task[None]: + entered = asyncio.Event() + restarting = self._track( + asyncio.create_task(announced_operation(self.client.start(), entered)) + ) + await entered.wait() + assert not restarting.done() + return restarting + + async def _claim_restart( + self, + restarting: asyncio.Task[None], + old_root: asyncio.subprocess.Process, + old_children: list[GatedProcess], + ) -> None: + await restarting + process = owned_rpc_process(self.client) + assert process is not old_root + self.roots.append(process) + await self.record_children() + self.live = True + await self.assert_reaped([old_root], old_children) + await second_prompt(self.client) + + @asynccontextmanager + async def held_first_spawn(self) -> AsyncIterator[HeldSpawn]: + await self.close() + barrier = SpawnBarrier(self.client.spawn_owned_process) + with patch.object(self.client, "spawn_owned_process", barrier.create): + starting = self._track(asyncio.create_task(self.client.start())) + try: + await barrier.created.wait() + assert barrier.process is not None + self.roots.append(barrier.process) + yield HeldSpawn( + starting=starting, + root=barrier.process, + children=await self.record_children(), + barrier=barrier, + ) + finally: + barrier.release.set() + + async def cancel_held_start_then_restart(self) -> None: + async with self.held_first_spawn() as spawn: + spawn.starting.cancel() + restarting = await self._queue_restart() + spawn.release() + await self._assert_finished(spawn.starting, asyncio.CancelledError) + await self._claim_restart(restarting, spawn.root, spawn.children) + + async def cancel_held_close_then_restart(self) -> None: + async with self.held_first_spawn() as spawn: + closing = await self._queue_close() + closing.cancel() + restarting = await self._queue_restart() + spawn.release() + await self._assert_finished(spawn.starting, OmpProcessError) + await self._assert_finished(closing, asyncio.CancelledError) + await self._claim_restart(restarting, spawn.root, spawn.children) + + async def take(self, step: Step) -> None: + match step: + case Step.PROMPT: + await self.hold_prompt() + case Step.CANCEL_QUEUED: + await self.cancel_queued() + case Step.SETTLE: + await self.settle() + case Step.CANCEL_OWNER: + await self.cancel_owner() + case Step.CLOSE: + await self.close() + case Step.CRASH: + await self.crash() + case Step.CANCEL_START_RESTART: + await self.cancel_held_start_then_restart() + case Step.CANCEL_CLOSE_RESTART: + await self.cancel_held_close_then_restart() + if self.live: + assert (await self.client.request("checkpoint"))["pid"] is not None + + async def assert_custody_survives(self) -> None: + """Settle leftover work, reap the tree, and prove a fresh turn still works.""" + if self.active is not None: + await self.take(Step.SETTLE) + await self.close() + assert all(task.done() for task in self.tasks) + await self.open() + await second_prompt(self.client) + await self.close() + + async def cleanup(self) -> None: + for task in self.tasks: + task.cancel() + await asyncio.gather(*self.tasks, return_exceptions=True) + await self.client.close() + + +@asynccontextmanager +async def rpc_lifecycle(*, detached: bool) -> AsyncIterator[RpcLifecycle]: + async with process_gate_session() as gate: + session = RpcLifecycle(gate, detached=detached) + try: + yield session + finally: + await session.cleanup() diff --git a/tests/integrations/omprpc/test_callback_retirement.py b/tests/integrations/omprpc/test_callback_retirement.py new file mode 100644 index 000000000..04dac3164 --- /dev/null +++ b/tests/integrations/omprpc/test_callback_retirement.py @@ -0,0 +1,119 @@ +"""Callbacks admitted during settlement share the original retirement deadline.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProcessError +from band.integrations.ompwire import ( + OmpCommandName, + OmpEventType, + OmpHostToolCall, + OmpToolResult, +) +from tests.integrations.omppeer import ( + LATE_CALLBACK_ID, + SETTLED_CALLBACK_ID, + SETTLED_CALLBACK_SCENARIO, +) +from tests.integrations.omprpc.cancellation import HeldCleanup +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + owned_rpc_process, +) +from tests.toolkit.clock import binary_clock_resolution +from tests.toolkit.timecontrol import controlled_clock + + +class TestCallbackRetirement: + @pytest.mark.parametrize( + "retire_second_on_first", + [False, True], + ids=["held-late-callback", "callback-turnover"], + ) + @pytest.mark.looptime(resolution=binary_clock_resolution()) + async def test_a_callback_admitted_during_the_join_cannot_escape_the_deadline( + self, client_factory: ClientFactory, retire_second_on_first: bool + ) -> None: + async with controlled_clock() as clock: + client = await clock.run_io(client_factory(SETTLED_CALLBACK_SCENARIO)) + grace = 1.0 + client.callback_exit_timeout_s = grace + process = owned_rpc_process(client) + observer = Observer() + cleanup = HeldCleanup(client) + joining: asyncio.Task[None] | None = None + first_started, first_retired, second_started = ( + asyncio.Event() for _ in range(3) + ) + release_first, release_second = (asyncio.Event() for _ in range(2)) + second_cancelled = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + if call.id == SETTLED_CALLBACK_ID: + task = asyncio.current_task() + assert task is not None + task.add_done_callback(lambda _done: first_retired.set()) + first_started.set() + await release_first.wait() + else: + assert call.id == LATE_CALLBACK_ID + second_started.set() + try: + await ( + first_retired.wait() + if retire_second_on_first + else release_second.wait() + ) + except asyncio.CancelledError: + second_cancelled.set() + raise + return OmpToolResult(content=[]) + + async with clock.background( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) as running: + try: + await first_started.wait() + await observer.wait_for(OmpEventType.AGENT_END) + frames = await checkpoint_frames(client) + assert any( + frame["type"] == OmpCommandName.GET_STATE for frame in frames + ) + await client.request("release") + await second_started.wait() + if retire_second_on_first: + clock.advance(grace) + release_first.set() + await first_retired.wait() + if not retire_second_on_first: + clock.advance(grace) + joining = asyncio.create_task(cleanup.wait_for_join(running)) + done, _ = await asyncio.wait( + (running, joining), return_when=asyncio.FIRST_COMPLETED + ) + assert joining in done + await joining + assert process.returncode is None + cleanup.release.set() + with pytest.raises(OmpProcessError): + await clock.run_io(running) + if not retire_second_on_first: + assert second_cancelled.is_set() + assert process.returncode is not None + finally: + release_first.set() + release_second.set() + cleanup.release.set() + if joining is not None: + joining.cancel() + await asyncio.gather(joining, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py new file mode 100644 index 000000000..fd0304bf6 --- /dev/null +++ b/tests/integrations/omprpc/test_callbacks.py @@ -0,0 +1,213 @@ +"""Host tool and UI callbacks run concurrently and propagate cancellation or errors.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.ompwire import ( + OmpContent, + OmpFrameType, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations.omppeer import ( + HOST_CANCELLED_ID, + UI_APPROVE, + UI_CANCELLED_ID, + UI_DENY, +) +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + owned_rpc_process, + second_prompt, +) + + +class TestHostToolAndUICallbacks: + async def test_host_tools_and_ui_run_concurrently_even_before_prompt_acceptance( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + first_started = asyncio.Event() + release_first = asyncio.Event() + calls: list[OmpHostToolCall] = [] + approvals: list[OmpUIRequest] = [] + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + calls.append(call) + if call.arguments["number"] == 1: + first_started.set() + await release_first.wait() + else: + await first_started.wait() + await client.request("checkpoint") + release_first.set() + return OmpToolResult( + content=[ + OmpContent(type="text", text=call.toolCallId), + OmpContent(type="image", data="aW1hZ2U=", mimeType="image/png"), + ], + details={"number": call.arguments["number"]}, + is_error=call.arguments["number"] == 2, + ) + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + approvals.append(request) + return OmpUIResponse(value=UI_APPROVE) + + async with asyncio.timeout(5): + await client.prompt( + "first", on_event=observer.on_event, on_tool=on_tool, on_ui=on_ui + ) + assert {call.id for call in calls} == {"host-1", "host-2"} + assert [(request.method, request.options) for request in approvals] == [ + ("select", [UI_APPROVE, UI_DENY]) + ] + frames = await checkpoint_frames(client) + result = next(frame for frame in frames if frame.get("id") == "host-2") + assert result == { + "type": OmpFrameType.HOST_TOOL_RESULT, + "id": "host-2", + "result": { + "content": [ + {"type": "text", "text": "call-2"}, + {"type": "image", "data": "aW1hZ2U=", "mimeType": "image/png"}, + ], + "details": {"number": 2}, + }, + "isError": True, + } + ui_response = next(frame for frame in frames if frame.get("id") == "approval-1") + assert ui_response == { + "type": OmpFrameType.UI_RESPONSE, + "id": "approval-1", + "value": UI_APPROVE, + } + await second_prompt(client) + + async def test_confirmed_ui_reply_reaches_the_peer( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + return OmpUIResponse(confirmed=True) + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + frames = await checkpoint_frames(client) + ui_response = next(frame for frame in frames if frame.get("id") == "approval-1") + assert ui_response == { + "type": OmpFrameType.UI_RESPONSE, + "id": "approval-1", + "confirmed": True, + } + + @pytest.mark.parametrize( + ("kind", "target"), + [("host", HOST_CANCELLED_ID), ("ui", UI_CANCELLED_ID)], + ) + async def test_native_cancel_targets_the_callback_and_keeps_the_next_turn_usable( + self, client_factory: ClientFactory, kind: str, target: str + ) -> None: + client = await client_factory(f"{kind}-cancel") + observer = Observer() + started = asyncio.Event() + cancelled = asyncio.Event() + forever = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpToolResult(content=[]) + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpUIResponse(value=UI_APPROVE) + + running = asyncio.create_task( + client.prompt( + "first", on_event=observer.on_event, on_tool=on_tool, on_ui=on_ui + ) + ) + try: + async with asyncio.timeout(5): + await started.wait() + await client.request(f"cancel-{kind}") + await cancelled.wait() + await running + frames = await checkpoint_frames(client) + assert not any(frame.get("id") == target for frame in frames) + await second_prompt(client) + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_ui_callback_failure_fails_closed_and_reaps_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("ui-error") + process = owned_rpc_process(client) + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + raise ValueError("Unoffered permission choice") + + observer = Observer() + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Unoffered permission choice"): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + assert process.returncode is not None + await client.start() + await second_prompt(client) + + async def test_tool_callback_exception_propagates_unchanged_and_reaps_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("callbacks") + process = owned_rpc_process(client) + observer = Observer() + sentinel = ValueError("Band delivery failed") + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + raise sentinel + + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Band delivery failed") as failed: + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + assert failed.value is sentinel + assert process.returncode is not None + await client.start() + await second_prompt(client) diff --git a/tests/integrations/omprpc/test_commands.py b/tests/integrations/omprpc/test_commands.py new file mode 100644 index 000000000..2c9efca4a --- /dev/null +++ b/tests/integrations/omprpc/test_commands.py @@ -0,0 +1,40 @@ +"""Command responses correlate independently and recover from errors or cancellation.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpRpcError +from tests.integrations.omprpc.kit import ClientFactory, second_prompt + + +class TestCommandCorrelation: + async def test_command_responses_are_correlated_out_of_order_and_errors_are_recoverable( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(5): + first, second = await asyncio.gather( + client.request("first"), client.request("second") + ) + assert (first, second) == ({"value": "first"}, {"value": "second"}) + with pytest.raises(OmpRpcError, match="Command failed") as failed: + await client.request("error") + assert (failed.value.command, failed.value.code) == ("error", "bad_command") + await second_prompt(client) + + async def test_cancelled_command_drops_only_its_waiter( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + pending = asyncio.create_task(client.request("hold")) + await client.request("checkpoint") + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert await client.request("second") == {"value": "second"} + await second_prompt(client) diff --git a/tests/integrations/omprpc/test_eof_descendants.py b/tests/integrations/omprpc/test_eof_descendants.py new file mode 100644 index 000000000..4d7072eca --- /dev/null +++ b/tests/integrations/omprpc/test_eof_descendants.py @@ -0,0 +1,77 @@ +"""Cleanup discovers descendants created by an owned helper after native stdin EOF.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from tests.integrations.omprpc.kit import ( + ProcessGate, + descendant_liveness, + hold_after_stdin_closed, + is_descendant, + job_member_pids, + process_is_alive, + session_ids, + start_eof_client, + unrelated_gate_child, +) +from tests.ostable import skip_unless_os + + +class TestEofDescendants: + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + async def test_close_reaps_a_detached_child_created_after_stdin_eof( + self, + process_gate: ProcessGate, + stdio: str, + ) -> None: + client, helper, root = await start_eof_client(process_gate, stdio=stdio) + async with unrelated_gate_child(process_gate) as unrelated: + try: + async with hold_after_stdin_closed(client) as closing: + late_child = await process_gate.connection() + assert root.returncode is not None + assert is_descendant(late_child.process, helper.process) + assert await descendant_liveness([helper, late_child]) == [ + True, + True, + ] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not closing.done() + assert await descendant_liveness([helper, late_child]) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_eof_late_child_is_its_own_session( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client, helper, root = await start_eof_client(process_gate, stdio="discard") + try: + async with hold_after_stdin_closed(client): + late_child = await process_gate.connection() + ids = session_ids([late_child.process.pid]) + assert ids == [late_child.process.pid] + assert ids[0] != root.pid + assert helper.process.pid != late_child.process.pid + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("nt",)]) + async def test_eof_late_child_is_a_job_member( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client, helper, _root = await start_eof_client(process_gate, stdio="discard") + try: + async with hold_after_stdin_closed(client): + late_child = await process_gate.connection() + members = job_member_pids(client) + assert helper.process.pid in members + assert late_child.process.pid in members + finally: + await client.close() diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py new file mode 100644 index 000000000..12df978c9 --- /dev/null +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -0,0 +1,188 @@ +"""Callback overload fails closed without blocking native control correlation.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +import pytest + +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.omprpc.errors import BACKLOG_EXCEEDED +from band.integrations.omprpc.prompt import EVENT_BACKLOG_LIMIT +from band.integrations.ompwire import ( + OmpEvent, + OmpEventType, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations.omppeer import COMPLETE, UI_APPROVE, burst_message_id +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + owned_rpc_process, + second_prompt, +) + + +@asynccontextmanager +async def held_event_prompt( + client: OmpRpcClient, +) -> AsyncIterator[ + tuple[asyncio.Task[None], asyncio.subprocess.Process, asyncio.Event] +]: + observer = Observer() + entered = asyncio.Event() + release = asyncio.Event() + stopped = asyncio.Event() + + async def on_event(event: OmpEvent) -> None: + entered.set() + try: + await release.wait() + finally: + stopped.set() + await observer.on_event(event) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + ) + try: + await entered.wait() + yield running, owned_rpc_process(client), stopped + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) + + +async def overflow_reaps_and_restarts( + client: OmpRpcClient, + running: asyncio.Task[None], + process: asyncio.subprocess.Process, + stopped: asyncio.Event, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.clear() + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc.client"): + burst = asyncio.create_task( + client.request("event-burst", count=EVENT_BACKLOG_LIMIT + 1) + ) + try: + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await running + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await burst + finally: + burst.cancel() + await asyncio.gather(burst, return_exceptions=True) + assert any( + record.levelno == logging.ERROR and BACKLOG_EXCEEDED in record.getMessage() + for record in caplog.records + ) + assert stopped.is_set() + assert process.returncode is not None + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + + +class TestEventBackpressure: + async def test_full_backlog_keeps_reentrant_control_live_and_preserves_event_order( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("event-backpressure") + process = owned_rpc_process(client) + observer = Observer() + entered = asyncio.Event() + release = asyncio.Event() + delivered = asyncio.Event() + burst_ids = [burst_message_id(number) for number in range(EVENT_BACKLOG_LIMIT)] + last_burst = burst_ids[-1] + + async def on_event(event: OmpEvent) -> None: + match event.type: + case OmpEventType.AGENT_START: + entered.set() + await release.wait() + await observer.on_event(event) + if event.messageId == last_burst: + delivered.set() + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + await client.request("checkpoint") + return OmpUIResponse(value=UI_APPROVE) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + ) + try: + async with asyncio.timeout(8): + await entered.wait() + await client.request("event-burst", count=EVENT_BACKLOG_LIMIT) + assert await client.request("event-ui") == {"value": UI_APPROVE} + assert observer.events == [] + assert not running.done() + assert process.returncode is None + + release.set() + await delivered.wait() + await client.request("event-result") + await running + + assert observer.outline == [ + OmpEventType.AGENT_START, + *([OmpEventType.MESSAGE_UPDATE] * EVENT_BACKLOG_LIMIT), + OmpEventType.MESSAGE_END, + OmpEventType.AGENT_END, + ] + assert observer.update_ids == burst_ids + assert observer.assistant_texts == [COMPLETE] + await second_prompt(client) + assert owned_rpc_process(client) is process + assert process.returncode is None + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) + + async def test_overload_before_result_fails_closed_and_restarts( + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture + ) -> None: + client = await client_factory("event-backpressure") + async with ( + asyncio.timeout(8), + held_event_prompt(client) as ( + running, + process, + stopped, + ), + ): + await overflow_reaps_and_restarts(client, running, process, stopped, caplog) + + async def test_overload_after_settlement_fails_closed_and_restarts( + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture + ) -> None: + client = await client_factory("event-backpressure") + async with ( + asyncio.timeout(8), + held_event_prompt(client) as ( + running, + process, + stopped, + ), + ): + await client.request("event-result") + await overflow_reaps_and_restarts(client, running, process, stopped, caplog) diff --git a/tests/integrations/omprpc/test_fail_closed.py b/tests/integrations/omprpc/test_fail_closed.py new file mode 100644 index 000000000..341a25161 --- /dev/null +++ b/tests/integrations/omprpc/test_fail_closed.py @@ -0,0 +1,91 @@ +"""Dispatcher fail-closed paths keep OMP from hanging on unmatched host frames.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProtocolError +from band.integrations.ompwire import OmpFrameType, OmpUIRequest, OmpUIResponse +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + prompt_task, +) + + +class TestFailClosedDispatch: + async def test_host_tool_without_an_active_prompt_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="without an active prompt"): + await client.request("orphan-tool") + + async def test_ui_cancel_without_target_id_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="missing targetId"): + await client.request("orphan-ui-cancel") + + async def test_interactive_ui_without_an_active_prompt_is_cancelled( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(5): + await client.request("orphan-ui") + frames = await checkpoint_frames(client) + orphan = next(frame for frame in frames if frame.get("id") == "ui-orphan") + assert orphan == { + "type": OmpFrameType.UI_RESPONSE, + "id": "ui-orphan", + "cancelled": True, + } + + async def test_ui_notifications_do_not_start_a_callback( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("ui-notify") + observer = Observer() + seen: list[OmpUIRequest] = [] + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + seen.append(request) + return await observer.on_ui(request) + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + assert seen == [] + + async def test_reused_callback_id_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("reused-callback") + observer = Observer() + async with asyncio.timeout(6): + with pytest.raises( + OmpProtocolError, match="callback request id was reused" + ): + await prompt_task(client, observer) + + async def test_request_rejects_id_and_type_overrides( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + with pytest.raises(ValueError, match="must not override"): + await client.request("checkpoint", id="stolen") + with pytest.raises(ValueError, match="must not override"): + await client.request("checkpoint", type="prompt") + async with asyncio.timeout(5): + checkpoint = await client.request("checkpoint") + assert "pid" in checkpoint diff --git a/tests/integrations/omprpc/test_injected_spawn.py b/tests/integrations/omprpc/test_injected_spawn.py new file mode 100644 index 000000000..2831aea53 --- /dev/null +++ b/tests/integrations/omprpc/test_injected_spawn.py @@ -0,0 +1,72 @@ +"""Injected creation preserves forced retirement and a usable replacement.""" + +from __future__ import annotations + +import asyncio +import sys + +import pytest + +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.ompwire import OmpEventType +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + FailingOwnedClient, + Observer, + owned_rpc_process, + prompt_task, + second_prompt, +) + + +class TestInjectedSpawnOwnership: + async def test_stubborn_injected_process_is_reaped_before_restart(self) -> None: + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", "stubborn"), + create_process=asyncio.create_subprocess_exec, + ) + process: asyncio.subprocess.Process | None = None + running: asyncio.Task[None] | None = None + try: + await client.start() + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + await observer.wait_for(OmpEventType.AGENT_END) + async with asyncio.timeout(10): + await client.close() + with pytest.raises(OmpProcessError): + await running + assert process.returncode is not None + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + finally: + # Release an uncontained direct child even when the regression fails. + if process is not None and process.returncode is None: + process.kill() + await process.wait() + if running is not None: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + await client.close() + + async def test_failed_owned_wrapper_reaps_before_restart(self) -> None: + client = FailingOwnedClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", "basic"), + create_process=asyncio.create_subprocess_exec, + ) + process: asyncio.subprocess.Process | None = None + try: + with pytest.raises(ValueError): + await client.start() + process = owned_rpc_process(client) + assert process.returncode is not None + client.fail_spawn = False + await client.start() + await second_prompt(client) + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + await client.close() diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py new file mode 100644 index 000000000..ccd7f0b52 --- /dev/null +++ b/tests/integrations/omprpc/test_interruption.py @@ -0,0 +1,210 @@ +"""Interrupted prompts and cancelled closes await callback cancellation and process reaping.""" + +from __future__ import annotations + +import asyncio +import sys + +import pytest +from pydantic import JsonValue + +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.ompwire import ( + OMP_FRAME_BYTES, + OmpEventType, + OmpHostToolCall, + OmpToolResult, +) +from tests.integrations.omppeer import STDIN_BACKPRESSURE +from tests.integrations.omprpc.backpressure import ( + PRESSURE_WRITE_COUNT, + BackpressureSpawn, +) +from tests.integrations.omprpc.cancellation import HeldCleanup +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ClientFactory, + Observer, + owned_rpc_process, + prompt_task, + queue_prompt, + second_prompt, + sent_prompt_messages, +) + + +class TestInterruptionAndReaping: + async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("yield") + process = owned_rpc_process(client) + observer = Observer() + active = prompt_task(client, observer) + queued: asyncio.Task[None] | None = None + try: + await observer.wait_for(OmpEventType.AGENT_END) + queued = await queue_prompt(client) + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + + assert not active.done() + assert process.returncode is None + assert await sent_prompt_messages(client) == ["first"] + await client.request("release") + async with asyncio.timeout(5): + await active + await second_prompt(client) + assert owned_rpc_process(client) is process + assert process.returncode is None + assert await sent_prompt_messages(client) == ["first", "second"] + finally: + tasks = [active, *([queued] if queued is not None else [])] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) + async def test_cancelled_prompt_aborts_reaps_and_can_restart( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + running.cancel() + async with asyncio.timeout(6): + with pytest.raises(asyncio.CancelledError): + await running + assert process.returncode is not None + async with asyncio.timeout(5): + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) + async def test_timed_out_prompt_aborts_reaps_and_can_restart( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + async with asyncio.timeout(6): + with pytest.raises(TimeoutError): + async with asyncio.timeout(None) as deadline: + deadline.reschedule(asyncio.get_running_loop().time()) + await running + assert deadline.expired() + assert process.returncode is not None + async with asyncio.timeout(5): + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_repeatedly_cancelled_close_still_waits_for_reap( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("stubborn") + process = owned_rpc_process(client) + cleanup = HeldCleanup(client) + closing = asyncio.create_task(client.close()) + try: + await cleanup.wait_for_join(closing) + for _ in range(3): + await cleanup.cancel_and_wait_for_rejoin(closing) + assert process.returncode is None + assert not cleanup.finished.is_set() + + cleanup.release.set() + async with asyncio.timeout(6): + with pytest.raises(asyncio.CancelledError): + await closing + await process.wait() + assert cleanup.finished.is_set() + assert process.returncode is not None + await client.close() + finally: + cleanup.release.set() + closing.cancel() + await asyncio.gather(closing, return_exceptions=True) + await client.close() + + async def test_cancelled_prompt_cancels_its_host_callback_before_reap( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("host-cancel") + process = owned_rpc_process(client) + observer = Observer() + started = asyncio.Event() + cancelled = asyncio.Event() + forever = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpToolResult(content=[]) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) + try: + async with asyncio.timeout(6): + await started.wait() + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + assert cancelled.is_set() + assert process.returncode is not None + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_close_finishes_while_a_stdin_drain_is_held(self) -> None: + spawn = BackpressureSpawn() + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", STDIN_BACKPRESSURE), + create_process=spawn.create, + ) + writes: list[asyncio.Task[dict[str, JsonValue]]] = [] + try: + await client.start() + process = owned_rpc_process(client) + await client.request("hold") + payload = "x" * (OMP_FRAME_BYTES // 2) + writes = [ + asyncio.create_task(client.request("pressure", data=payload)) + for _ in range(PRESSURE_WRITE_COUNT) + ] + async with asyncio.timeout(5): + await spawn.paused.wait() + await client.close() + assert process.returncode is not None + outcomes = await asyncio.gather(*writes, return_exceptions=True) + assert all(isinstance(outcome, OmpProcessError) for outcome in outcomes) + finally: + for writing in writes: + writing.cancel() + await asyncio.gather(*writes, return_exceptions=True) + await client.close() diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py new file mode 100644 index 000000000..0a851c2eb --- /dev/null +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -0,0 +1,63 @@ +"""Owned process trees and admitted turns survive generated host interleavings.""" + +from __future__ import annotations + +import asyncio + +import pytest +from hypothesis import Phase, example, given, settings + +from tests.integrations.omprpc.lifecycle import ( + STEP_GUARD_S, + Step, + rpc_lifecycle, + sequences, +) + + +async def play(steps: list[Step], *, detached: bool) -> None: + async with rpc_lifecycle(detached=detached) as session: + for step in steps: + async with asyncio.timeout(STEP_GUARD_S): + await session.take(step) + async with asyncio.timeout(STEP_GUARD_S): + await session.assert_custody_survives() + + +class TestRpcLifecycleModel: + @pytest.mark.parametrize("detached", [False, True], ids=["session", "detached"]) + # The full Hypothesis workload can exceed two minutes on Windows. + # Each generated step retains its separate, shorter hang guard. + @pytest.mark.timeout(600) + @settings( + max_examples=20, + deadline=None, + phases=(Phase.explicit, Phase.reuse, Phase.generate), + ) + @given(steps=sequences) + @example(steps=[Step.PROMPT, Step.CANCEL_QUEUED, Step.SETTLE]) + @example(steps=[Step.PROMPT, Step.CRASH, Step.PROMPT, Step.CANCEL_OWNER]) + @example(steps=[Step.CANCEL_START_RESTART, Step.PROMPT, Step.CLOSE]) + @example(steps=[Step.CANCEL_CLOSE_RESTART, Step.PROMPT, Step.SETTLE]) + @example( + steps=[ + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_CLOSE_RESTART, + ] + ) + @example(steps=[Step.PROMPT, Step.PROMPT, Step.CANCEL_OWNER, Step.CANCEL_OWNER]) + @example( + steps=[ + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_OWNER, + ] + ) + def test_generated_operations_preserve_custody_and_settlement( + self, steps: list[Step], detached: bool + ) -> None: + asyncio.run(play(steps, detached=detached)) diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py new file mode 100644 index 000000000..e80f7d270 --- /dev/null +++ b/tests/integrations/omprpc/test_process_tree.py @@ -0,0 +1,321 @@ +"""Cancelled cleanup reaps descendants without affecting unrelated processes.""" + +from __future__ import annotations + +import asyncio + +import psutil +import pytest + +from band.integrations.omprpc import OmpProcessError +from band.integrations.omprpc.client import PROCESS_EXIT_TIMEOUT_S +from band.integrations.omprpc.scope import OmpProcessScope +from tests.integrations.omprpc.kit import ( + ClientFactory, + DeferredSessionSignal, + ProcessGate, + RecycledEnvironment, + RecycledGroup, + connect_descendant_tree, + crash_rpc_root, + descendant_client, + descendant_liveness, + job_member_pids, + owned_rpc_process, + process_is_alive, + refuse_containment, + session_ids, + unrelated_gate_child, + wait_for_descendants, +) +from tests.ostable import skip_unless_os + + +class TestProcessTreeCleanup: + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_reused_group_during_capture_cannot_adopt_an_unrelated_process( + self, + client_factory: ClientFactory, + process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, + accepted_on: tuple[str, ...], + ) -> None: + skip_unless_os(accepted_on) + client = await client_factory("basic") + root = owned_rpc_process(client) + await client.close() + assert root.returncode is not None + scope = OmpProcessScope() + async with unrelated_gate_child(process_gate) as unrelated: + boundary = RecycledGroup(root, unrelated) + with monkeypatch.context() as patch: + boundary.install(patch) + scope.capture(root) + assert boundary.reused + assert unrelated.process not in scope.processes + assert await asyncio.to_thread(process_is_alive, unrelated.process) + + async def test_close_reaps_detached_descendants_without_affecting_other_owners( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate, detached=True) + other = descendant_client(process_gate, detached=True) + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + await other.start() + other_descendants = await wait_for_descendants(process_gate) + other_root = owned_rpc_process(other) + + await client.close() + assert root.returncode is not None + assert await descendant_liveness(descendants) == [False, False] + assert other_root.returncode is None + assert await descendant_liveness(other_descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert (await other.request("get_state"))["isSettled"] is True + await other.close() + assert await descendant_liveness(other_descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + await other.close() + + async def test_crash_then_close_reaps_detached_descendants_without_affecting_other_owners( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate, detached=True) + other = descendant_client(process_gate, detached=True) + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + await other.start() + other_descendants = await wait_for_descendants(process_gate) + other_root = owned_rpc_process(other) + await crash_rpc_root(client) + + await client.close() + assert root.returncode is not None + assert await descendant_liveness(descendants) == [False, False] + assert other_root.returncode is None + assert await descendant_liveness(other_descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await other.close() + assert await descendant_liveness(other_descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + await other.close() + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_detached_descendants_are_their_own_sessions( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client = descendant_client(process_gate, detached=True) + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + ids = session_ids([child.process.pid for child in descendants]) + assert ids == [child.process.pid for child in descendants] + assert all(session != root.pid for session in ids) + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("nt",)]) + async def test_detached_descendants_are_job_members( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + """Windows identity complementary to POSIX session ids, not a nested-job claim.""" + skip_unless_os(accepted_on) + client = descendant_client(process_gate, detached=True) + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + members = job_member_pids(client) + assert {child.process.pid for child in descendants} <= members + finally: + await client.close() + + async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity( + self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = descendant_client(process_gate, detached=True) + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + boundary = RecycledEnvironment(unrelated, root.pid) + with monkeypatch.context() as patch: + boundary.install(patch) + await client.close() + assert boundary.recycled + assert await descendant_liveness(descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + + async def test_close_reaps_owned_descendants_after_the_root_has_crashed( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate) + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + assert await descendant_liveness(descendants) == [True, True] + await crash_rpc_root(client) + await client.close() + assert await descendant_liveness(descendants) == [False, False] + finally: + await client.close() + + async def test_failed_containment_cleanup_is_retried_before_restart( + self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = descendant_client(process_gate) + try: + await client.start() + old_root = owned_rpc_process(client) + descendants = await wait_for_descendants(process_gate) + with monkeypatch.context() as patch: + refuse_containment(patch) + await crash_rpc_root(client) + with pytest.raises(PermissionError, match="not signalable"): + await client.close() + assert await descendant_liveness(descendants) == [True, True] + + await client.start() + assert owned_rpc_process(client) is not old_root + assert await descendant_liveness(descendants) == [False, False] + await wait_for_descendants(process_gate) + await client.close() + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + @pytest.mark.parametrize( + "hide_after_signal", [False, True], ids=["visible", "hidden-after-signal"] + ) + async def test_unsettled_session_members_keep_close_retryable( + self, + process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, + accepted_on: tuple[str, ...], + hide_after_signal: bool, + ) -> None: + skip_unless_os(accepted_on) + client = descendant_client(process_gate) + try: + await client.start() + old_root = owned_rpc_process(client) + descendants = await wait_for_descendants(process_gate) + boundary = DeferredSessionSignal( + old_root.pid, descendants, hide_after_signal=hide_after_signal + ) + with monkeypatch.context() as patch: + boundary.install(patch) + client.process_exit_timeout_s = 0 + await crash_rpc_root(client) + for _ in range(2): + with pytest.raises(OmpProcessError, match="session members"): + await client.close() + assert await descendant_liveness(descendants) == [True, True] + client.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S + + await client.start() + assert owned_rpc_process(client) is not old_root + assert await descendant_liveness(descendants) == [False, False] + await wait_for_descendants(process_gate) + await client.close() + finally: + await client.close() + + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + async def test_cancelled_startup_cleanup_reaps_descendants_but_not_unrelated_processes( + self, + process_gate: ProcessGate, + stdio: str, + ) -> None: + tree = await connect_descendant_tree( + process_gate, scenario="no-ready-descendants", stdio=stdio + ) + async with unrelated_gate_child(process_gate) as unrelated: + try: + root_process = await asyncio.to_thread(psutil.Process, tree.root.pid) + closing = tree.starting + assert closing.cancel() + await tree.cleanup.wait_for_join(closing) + for _ in range(3): + await tree.cleanup.cancel_and_wait_for_rejoin(closing) + assert tree.root.returncode is None + assert await asyncio.to_thread(process_is_alive, root_process) + assert await descendant_liveness(tree.descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not tree.cleanup.finished.is_set() + + tree.cleanup.release.set() + async with asyncio.timeout(10): + with pytest.raises(asyncio.CancelledError): + await closing + await tree.root.wait() + assert tree.cleanup.finished.is_set() + assert tree.root.returncode is not None + assert await descendant_liveness(tree.descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await tree.client.close() + finally: + tree.cleanup.release.set() + tree.starting.cancel() + await asyncio.gather(tree.starting, return_exceptions=True) + await tree.client.close() + + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + async def test_cancelled_close_cleanup_reaps_descendants_but_not_unrelated_processes( + self, + process_gate: ProcessGate, + stdio: str, + ) -> None: + tree = await connect_descendant_tree( + process_gate, scenario="descendants", stdio=stdio + ) + closing: asyncio.Task[None] | None = None + async with unrelated_gate_child(process_gate) as unrelated: + try: + async with asyncio.timeout(5): + await tree.starting + root_process = await asyncio.to_thread(psutil.Process, tree.root.pid) + closing = asyncio.create_task(tree.client.close()) + await tree.cleanup.wait_for_join(closing) + await tree.cleanup.cancel_and_wait_for_rejoin(closing) + for _ in range(3): + await tree.cleanup.cancel_and_wait_for_rejoin(closing) + assert tree.root.returncode is None + assert await asyncio.to_thread(process_is_alive, root_process) + assert await descendant_liveness(tree.descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not tree.cleanup.finished.is_set() + + tree.cleanup.release.set() + async with asyncio.timeout(10): + with pytest.raises(asyncio.CancelledError): + await closing + await tree.root.wait() + assert tree.cleanup.finished.is_set() + assert tree.root.returncode is not None + assert await descendant_liveness(tree.descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await tree.client.close() + finally: + tree.cleanup.release.set() + tree.starting.cancel() + tasks = [tree.starting, *([closing] if closing is not None else [])] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await tree.client.close() diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py new file mode 100644 index 000000000..d1e43d34f --- /dev/null +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -0,0 +1,218 @@ +"""Prompt narration and results preserve typed content and report failures.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpPromptError, OmpProtocolError, OmpRpcError +from band.integrations.ompwire import ( + OMP_NATIVE_PROTOCOL_VERSION, + OmpCommandName, + OmpEvent, + OmpEventType, + OmpState, +) +from tests.integrations.omppeer import ( + ADMITTED_PROMPT_ERROR, + COMPLETE, + LATE_PROMPT_ERROR, +) +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + owned_rpc_process, + prompt_task, + second_prompt, +) + + +class TestPromptNarrationAndResults: + async def test_prompt_delivers_typed_narration_and_ignores_other_events( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + assert observer.outline == [ + OmpEventType.AGENT_START, + OmpEventType.TOOL_START, + OmpEventType.TOOL_END, + OmpEventType.MESSAGE_END, + OmpEventType.AGENT_END, + ] + tool_start, tool_end, message_end = observer.events[1:4] + assert (tool_start.toolCallId, tool_start.toolName, tool_start.args) == ( + "native-call", + "read", + {"path": "note.txt"}, + ) + assert tool_end.result is not None + assert tool_end.result.content[0].text == "note" + assert tool_end.result.details == {"lines": 1} + assert tool_end.isError is False + message = message_end.message + assert message is not None + assert message.role == "assistant" + assert isinstance(message.content, list) + assert [ + (content.type, content.text, content.thinking) + for content in message.content + ] == [ + ("thinking", None, "Considering the result"), + ("text", COMPLETE, None), + ] + assert message.usage is not None + assert ( + message.usage.input, + message.usage.output, + message.usage.totalTokens, + ) == ( + 7, + 3, + 13, + ) + assert message.usage.cost is not None + assert message.usage.cost.total == 0.33 + state = OmpState.model_validate(await client.request("get_state")) + assert (state.sessionId, state.sessionFile, state.isSettled) == ( + "session-local", + "/local/session.jsonl", + True, + ) + assert state.model is not None + assert (state.model.provider, state.model.id) == ("test-provider", "test-model") + await second_prompt(client) + + async def test_rejected_prompt_does_not_wait_for_a_result_that_is_not_owed( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("rejected") + async with asyncio.timeout(5): + with pytest.raises(OmpRpcError, match="Prompt rejected") as rejected: + await prompt_task(client, Observer()) + assert (rejected.value.command, rejected.value.code) == ("prompt", "rejected") + await second_prompt(client) + + async def test_command_error_after_admission_keeps_the_process_usable( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(LATE_PROMPT_ERROR) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + await client.request("release") + with pytest.raises(OmpRpcError) as failed: + await running + assert failed.value.command == OmpCommandName.PROMPT + assert process.returncode is None + await second_prompt(client) + assert owned_rpc_process(client) is process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_a_late_response_for_the_wrong_command_retires_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(LATE_PROMPT_ERROR) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + with pytest.raises(OmpProtocolError): + await client.request("release", failed_command=OmpCommandName.GET_STATE) + with pytest.raises(OmpProtocolError): + await running + assert process.returncode is not None + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_a_late_prompt_error_does_not_reap_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory(ADMITTED_PROMPT_ERROR) + process = owned_rpc_process(client) + async with asyncio.timeout(5): + with pytest.raises(OmpRpcError) as late: + await prompt_task(client, Observer()) + assert (late.value.command, late.value.code) == ( + OmpCommandName.PROMPT, + "post_admit", + ) + assert process.returncode is None + await second_prompt(client) + + @pytest.mark.parametrize( + ("scenario", "match"), + [ + ("initial-error", "Initial provider failed"), + ("initial-aborted", "OMP prompt aborted"), + ], + ) + async def test_failed_initial_result_is_reported( + self, client_factory: ClientFactory, scenario: str, match: str + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + async with asyncio.timeout(5): + with pytest.raises(OmpPromptError, match=match): + await prompt_task(client, Observer()) + assert process.returncode is None + await second_prompt(client) + + async def test_narration_failure_does_not_leave_the_turn_or_process_waiting( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + process = owned_rpc_process(client) + observer = Observer() + + async def on_event(event: OmpEvent) -> None: + raise ValueError("Narration failed") + + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Narration failed"): + await client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + assert process.returncode is not None + + async def test_native_v2_preserves_large_inline_image_frames( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("large-image") + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + message = next( + event.message + for event in observer.events + if event.type == OmpEventType.MESSAGE_END + ) + assert message is not None and isinstance(message.content, list) + image = next(content for content in message.content if content.type == "image") + assert image.mimeType == "image/png" + assert image.data == "YWJj" * 300_000 + negotiations = [ + frame + for frame in await checkpoint_frames(client) + if frame["type"] == OmpCommandName.NEGOTIATE_PROTOCOL + ] + assert len(negotiations) == 1 + assert negotiations[0]["protocolVersion"] == OMP_NATIVE_PROTOCOL_VERSION diff --git a/tests/integrations/omprpc/test_startup.py b/tests/integrations/omprpc/test_startup.py new file mode 100644 index 000000000..6233c80be --- /dev/null +++ b/tests/integrations/omprpc/test_startup.py @@ -0,0 +1,95 @@ +"""Startup requires native protocol negotiation and reaps failed launches.""" + +from __future__ import annotations + +import asyncio +import logging +import sys + +import pytest + +from band.integrations.omprpc import ( + OmpProcessError, + OmpProtocolError, + OmpRpcClient, + OmpRpcError, +) +from tests.integrations.omppeer import OVERSIZE_AFTER_READY +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ClientFactory, + SpawnBarrier, + owned_rpc_process, +) + + +class TestStartupNegotiation: + async def test_cancelled_startup_reaps_a_process_that_never_announces_ready( + self, + ) -> None: + barrier = SpawnBarrier(asyncio.create_subprocess_exec) + client = OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + "no-ready", + ), + create_process=barrier.create, + ) + starting = asyncio.create_task(client.start()) + try: + async with asyncio.timeout(5): + await barrier.created.wait() + assert barrier.process is not None + assert starting.cancel() + barrier.release.set() + with pytest.raises(asyncio.CancelledError): + await starting + assert barrier.process.returncode is not None + finally: + barrier.release.set() + await client.close() + + async def test_missing_executable_is_a_process_failure(self) -> None: + client = OmpRpcClient(("definitely-not-an-omp-executable-xyz",)) + try: + with pytest.raises(OmpProcessError, match="start"): + await client.start() + finally: + await client.close() + + @pytest.mark.parametrize( + ("scenario", "error"), + [ + ("no-v2", OmpProtocolError), + ("bad-negotiation", OmpProtocolError), + ("negotiation-rejected", OmpRpcError), + ], + ) + async def test_start_requires_successful_native_v2_negotiation( + self, client_factory: ClientFactory, scenario: str, error: type[RuntimeError] + ) -> None: + async with asyncio.timeout(6): + with pytest.raises(error, match="protocol"): + await client_factory(scenario) + + async def test_negotiated_frame_limit_applies_before_the_next_line(self) -> None: + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", OVERSIZE_AFTER_READY) + ) + try: + with pytest.raises(OmpProtocolError): + await client.start() + assert owned_rpc_process(client).returncode is not None + finally: + await client.close() + + async def test_graceful_close_does_not_log_stdout_eof_as_a_crash( + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture + ) -> None: + client = await client_factory("basic") + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc"): + await client.close() + assert caplog.records == [] diff --git a/tests/integrations/omprpc/test_transport_failures.py b/tests/integrations/omprpc/test_transport_failures.py new file mode 100644 index 000000000..cfc1065a9 --- /dev/null +++ b/tests/integrations/omprpc/test_transport_failures.py @@ -0,0 +1,148 @@ +"""Transport and protocol failures fail waiters and reap their subprocesses.""" + +from __future__ import annotations + +import asyncio +import logging + +import pytest +from pydantic import JsonValue + +from band.integrations.codex.winspawn import await_returncode +from band.integrations.omprpc import OmpProcessError, OmpProtocolError +from band.integrations.ompwire import OmpEventType +from tests.integrations.omppeer import FRAME_PAYLOAD_MARK +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + ProcessGate, + descendant_client, + descendant_liveness, + owned_rpc_process, + prompt_task, + unrelated_gate_child, + wait_for_descendants, +) + + +class TestTransportAndProtocolFailures: + @pytest.mark.parametrize( + ("scenario", "error", "match"), + [ + ("eof", OmpProcessError, "EOF"), + ("bad-json", OmpProtocolError, "protocol"), + ("bad-event", OmpProtocolError, "protocol"), + ("wrong-command", OmpProtocolError, "command"), + ("missing-state", OmpProtocolError, "protocol"), + ("bad-acceptance", OmpProtocolError, "protocol"), + ], + ) + async def test_eof_and_protocol_failure_cannot_leave_a_prompt_waiting( + self, + client_factory: ClientFactory, + scenario: str, + error: type[RuntimeError], + match: str, + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + async with asyncio.timeout(6): + with pytest.raises(error, match=match): + await prompt_task(client, Observer()) + assert process.returncode is not None + + async def test_protocol_failure_logs_omit_the_frame_payload( + self, + client_factory: ClientFactory, + caplog: pytest.LogCaptureFixture, + ) -> None: + client = await client_factory("bad-json") + caplog.clear() + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc.client"): + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="protocol"): + await prompt_task(client, Observer()) + assert FRAME_PAYLOAD_MARK not in caplog.text + + async def test_eof_fails_a_command_waiter_and_closes_the_transport( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + process = owned_rpc_process(client) + async with asyncio.timeout(6): + with pytest.raises(OmpProcessError, match="EOF"): + await client.request("crash") + await client.close() + assert process.returncode is not None + + async def test_observed_root_exit_retires_an_inflight_prompt_and_owned_children( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate, stdio="inherit", detached=True) + observer = Observer() + running: asyncio.Task[None] | None = None + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + process = owned_rpc_process(client) + running = prompt_task(client, observer) + await observer.wait_for(OmpEventType.AGENT_END) + crashing = asyncio.create_task(client.request("crash")) + assert await await_returncode(process, float("inf")) + assert await descendant_liveness(descendants) == [True, True] + + with pytest.raises(OmpProcessError): + await client.request("get_state") + async with asyncio.timeout(5): + with pytest.raises(OmpProcessError): + await running + with pytest.raises(OmpProcessError): + await crashing + assert await descendant_liveness(descendants) == [False, False] + assert await descendant_liveness([unrelated]) == [True] + finally: + await client.close() + tasks = [task for task in (running, crashing) if task is not None] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def test_stderr_is_drained_without_newlines( + self, + client_factory: ClientFactory, + caplog: pytest.LogCaptureFixture, + ) -> None: + client = await client_factory("basic") + with caplog.at_level(logging.DEBUG, logger="band.integrations.omprpc.client"): + async with asyncio.timeout(5): + assert await client.request("stderr") == {"drained": True} + leaked = any( + FRAME_PAYLOAD_MARK in record.getMessage() for record in caplog.records + ) + caplog.clear() + assert not leaked + + @pytest.mark.parametrize( + "scenario", + [ + "chunk-overflow", + "chunk-order", + "chunk-interrupted", + "chunk-invalid-data", + "chunk-length", + "chunk-eof", + "frame-overflow", + ], + ) + async def test_invalid_or_oversized_native_chunks_fail_and_reap( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="protocol"): + await prompt_task(client, Observer()) + assert process.returncode is not None diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py new file mode 100644 index 000000000..d9545c57c --- /dev/null +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -0,0 +1,168 @@ +"""Turns await matching results and settlement across background work.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProcessError, OmpPromptError +from band.integrations.omprpc.errors import CALLBACKS_UNFINISHED +from band.integrations.ompwire import ( + OmpEventType, + OmpHostToolCall, + OmpPromptStatus, + OmpState, + OmpToolResult, +) +from tests.integrations.omppeer import ( + BACKGROUND_COMPLETE, + RESUMED_PROVIDER_FAILED, + SETTLED_CALLBACK_SCENARIO, + YIELDED, +) +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + owned_rpc_process, + prompt_task, + second_prompt, +) + + +class TestTurnSettlement: + @pytest.mark.parametrize("scenario", ["yield", "stale", "false-result-settled"]) + async def test_initial_yield_never_finishes_a_background_turn( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + running = prompt_task(client, observer) + try: + first_end = await observer.wait_for(OmpEventType.AGENT_END) + assert first_end.awaitingAsyncWork is True + await client.request("checkpoint") + assert not running.done() + await client.request("release") + async with asyncio.timeout(5): + await running + assert observer.assistant_texts == [YIELDED, BACKGROUND_COMPLETE] + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize( + ("scenario", "error"), + [ + ("resume-error", RESUMED_PROVIDER_FAILED), + ("resume-error-compacted", RESUMED_PROVIDER_FAILED), + ( + "resume-aborted", + f"OMP continuation {OmpPromptStatus.ABORTED}", + ), + ], + ) + async def test_resumed_failures_override_a_successful_initial_prompt_result( + self, client_factory: ClientFactory, scenario: str, error: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + await client.request("release") + async with asyncio.timeout(5): + with pytest.raises(OmpPromptError, match=error): + await running + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_settlement_before_matching_result_waits_for_both( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("settlement-before-result") + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + await client.request("checkpoint") + assert not running.done() + await client.request("release") + async with asyncio.timeout(5): + await running + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize( + "scenario", ["state-settles", "settles-during-state", "local"] + ) + async def test_native_state_snapshot_and_subscribed_settlement_do_not_race( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + assert OmpState.model_validate(await client.request("get_state")).isSettled + + async def test_settlement_does_not_succeed_while_a_callback_is_still_running( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(SETTLED_CALLBACK_SCENARIO) + client.callback_exit_timeout_s = 0 + process = owned_rpc_process(client) + held = asyncio.Event() + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + await held.wait() + return OmpToolResult(content=[]) + + observer = Observer() + with pytest.raises(OmpProcessError, match=CALLBACKS_UNFINISHED): + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + held.set() + assert process.returncode is not None + await client.start() + await second_prompt(client) + + async def test_settlement_succeeds_after_the_running_callback_finishes( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(SETTLED_CALLBACK_SCENARIO) + started = asyncio.Event() + release = asyncio.Event() + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + started.set() + await release.wait() + return OmpToolResult(content=[]) + + observer = Observer() + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) + try: + await started.wait() + release.set() + async with asyncio.timeout(5): + await running + assert owned_rpc_process(client).returncode is None + await second_prompt(client) + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_ui_answers.py b/tests/integrations/omprpc/test_ui_answers.py new file mode 100644 index 000000000..a5616ed70 --- /dev/null +++ b/tests/integrations/omprpc/test_ui_answers.py @@ -0,0 +1,66 @@ +"""Structured UI answers reach the native peer without becoming cancellation.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import JsonValue + +from band.integrations.ompwire import ( + OmpEventType, + OmpFrameType, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations.omprpc.kit import ClientFactory, Observer + + +class TestStructuredUIAnswers: + @pytest.mark.parametrize( + "answer", + [ + {"id": "database-choice", "selectedOptions": []}, + { + "id": "database-choice", + "selectedOptions": [], + "customInput": "Use the existing database", + }, + ], + ids=["empty-selection", "custom-input"], + ) + async def test_answer_variant_survives_callback_to_wire( + self, client_factory: ClientFactory, answer: dict[str, JsonValue] + ) -> None: + client = await client_factory("event-backpressure") + observer = Observer() + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + assert request.method == "ask" + assert request.questions[0].id == answer["id"] + return OmpUIResponse.model_validate({"answers": [answer]}) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + ) + try: + async with asyncio.timeout(8): + await observer.wait_for(OmpEventType.AGENT_START) + receipt = await client.request("ask-ui") + assert receipt == { + "answer": { + "type": OmpFrameType.UI_RESPONSE, + "id": "structured-question", + "answers": [answer], + } + } + await client.request("event-result") + await running + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_ui_responses.py b/tests/integrations/omprpc/test_ui_responses.py new file mode 100644 index 000000000..2def5fd95 --- /dev/null +++ b/tests/integrations/omprpc/test_ui_responses.py @@ -0,0 +1,65 @@ +"""Preserve native UI answer variants through the public prompt callback.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import JsonValue + +from band.integrations.ompwire import OmpFrameType, OmpUIRequest, OmpUIResponse +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + second_prompt, +) + + +class TestUIResponseVariants: + @pytest.mark.parametrize( + ("response", "answer"), + [ + pytest.param(OmpUIResponse(value=""), {"value": ""}, id="empty-text"), + pytest.param( + OmpUIResponse(confirmed=False), + {"confirmed": False}, + id="declined-confirmation", + ), + pytest.param( + OmpUIResponse(cancelled=False), {}, id="unanswered-not-cancelled" + ), + pytest.param( + OmpUIResponse(cancelled=True), + {"cancelled": True}, + id="explicit-cancellation", + ), + ], + ) + async def test_callback_answer_reaches_the_peer_without_other_variants( + self, + client_factory: ClientFactory, + response: OmpUIResponse, + answer: dict[str, JsonValue], + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + return response + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + frames = await checkpoint_frames(client) + reply = next(frame for frame in frames if frame.get("id") == "approval-1") + assert reply == { + "type": OmpFrameType.UI_RESPONSE, + "id": "approval-1", + **answer, + } + await second_prompt(client) diff --git a/tests/integrations/test_ompwire.py b/tests/integrations/test_ompwire.py new file mode 100644 index 000000000..012374774 --- /dev/null +++ b/tests/integrations/test_ompwire.py @@ -0,0 +1,74 @@ +"""Structured ask frames preserve native metadata and mutually exclusive answers.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from band.integrations.ompwire import OmpUIRequest, OmpUIResponse + + +class TestOmpAskWire: + def test_structured_questions_preserve_native_metadata(self) -> None: + question = { + "id": "database-choice", + "question": "Which database should this workspace use?", + "header": "Database", + "options": [ + { + "label": "SQLite", + "description": "A local database", + "preview": "CREATE TABLE workspace (id INTEGER);", + } + ], + "multi": True, + "recommended": 0, + } + request = OmpUIRequest.model_validate( + {"id": "ask-1", "method": "ask", "questions": [question]} + ) + + assert request.model_dump(mode="json", exclude_none=True)["questions"] == [ + question + ] + + @pytest.mark.parametrize("recommended", [-1, 0.5, 8]) + def test_numeric_recommendations_do_not_break_the_reader( + self, recommended: float + ) -> None: + request = OmpUIRequest.model_validate( + { + "id": "ask-2", + "method": "ask", + "questions": [ + { + "id": "question-1", + "question": "Pick a database", + "options": [{"label": "SQLite", "preview": "schema preview"}], + "recommended": recommended, + } + ], + } + ) + assert request.questions[0].recommended == recommended + assert request.questions[0].options[0].preview == "schema preview" + + def test_structured_answers_preserve_an_empty_native_selection(self) -> None: + answers = [{"id": "question-2", "selectedOptions": []}] + response = OmpUIResponse.model_validate({"answers": answers}) + + assert response.model_dump(mode="json", exclude_none=True)["answers"] == answers + + @pytest.mark.parametrize( + "other_answer", [{"value": "SQLite"}, {"confirmed": True}, {"cancelled": True}] + ) + def test_structured_answers_cannot_be_mixed_with_another_ui_variant( + self, other_answer: dict[str, str | bool] + ) -> None: + with pytest.raises(ValidationError, match="only one answer"): + OmpUIResponse.model_validate( + { + "answers": [{"id": "question-3", "selectedOptions": []}], + **other_answer, + } + ) diff --git a/tests/ostable.py b/tests/ostable.py new file mode 100644 index 000000000..ae5860324 --- /dev/null +++ b/tests/ostable.py @@ -0,0 +1,12 @@ +"""OS-table rows that run everywhere and skip only on the named platforms.""" + +from __future__ import annotations + +import os + +import pytest + + +def skip_unless_os(accepted_on: tuple[str, ...]) -> None: + if os.name not in accepted_on: + pytest.skip(f"this row runs on {', '.join(accepted_on)}") diff --git a/tests/toolkit/timecontrol.py b/tests/toolkit/timecontrol.py index 8e7fa4c6d..3212063cd 100644 --- a/tests/toolkit/timecontrol.py +++ b/tests/toolkit/timecontrol.py @@ -3,8 +3,10 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +import threading +from collections.abc import AsyncIterator, Awaitable, Coroutine +from contextlib import AsyncExitStack, asynccontextmanager +from typing import Any, TypeVar from looptime import LoopTimeEventLoop @@ -12,12 +14,17 @@ from tests.toolkit.clock import binary_clock_resolution from tests.toolkit.tasks import background +Result = TypeVar("Result") + +REAL_IO_STEP_S = 0.01 + class ControlledClock: """Keep time stationary until the test crosses a deliberate boundary.""" - def __init__(self, loop: LoopTimeEventLoop) -> None: + def __init__(self, loop: LoopTimeEventLoop, stationary: asyncio.Task[None]) -> None: self.loop = loop + self._stationary = stationary def advance(self, seconds: float) -> None: self.loop.setup_looptime( @@ -29,6 +36,48 @@ async def elapse(self, seconds: float) -> None: self.advance(seconds) await spin() + async def _pause(self) -> None: + self._stationary.cancel() + await asyncio.gather(self._stationary, return_exceptions=True) + + async def run_io(self, operation: Awaitable[Result]) -> Result: + """Let real subprocess readiness and retirement make timed progress.""" + # Edge-triggered I/O cannot block on a second infinite selector poll + # before the loop services the completion returned by the first one. + self.loop.setup_looptime( + start=self.loop.time(), + resolution=binary_clock_resolution(), + idle_step=REAL_IO_STEP_S, + ) + running = asyncio.ensure_future(operation) + release = threading.Event() + # Looptime follows real time while an executor operation is pending. + # Hold that mode until the whole I/O operation, including cleanup, ends. + waiting = asyncio.create_task(asyncio.to_thread(release.wait)) + try: + await self._pause() + return await running + finally: + try: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + finally: + release.set() + await waiting + self._stationary = asyncio.create_task(keep_time_stationary()) + + @asynccontextmanager + async def background( + self, operation: Coroutine[Any, Any, Result] | asyncio.Task[Result] + ) -> AsyncIterator[asyncio.Task[Result]]: + """Join a native operation with real I/O progress on exceptional exits.""" + owner = AsyncExitStack() + try: + task = await owner.enter_async_context(background(operation)) + yield task + finally: + await self.run_io(owner.aclose()) + async def keep_time_stationary() -> None: # A ready task prevents looptime from jumping to a timer during socket I/O. @@ -40,5 +89,9 @@ async def keep_time_stationary() -> None: async def controlled_clock() -> AsyncIterator[ControlledClock]: loop = asyncio.get_running_loop() assert isinstance(loop, LoopTimeEventLoop) and loop.looptime_on - async with background(keep_time_stationary()): - yield ControlledClock(loop) + async with background(keep_time_stationary()) as stationary: + clock = ControlledClock(loop, stationary) + try: + yield clock + finally: + await clock._pause() diff --git a/uv.lock b/uv.lock index 8684ecbef..12ff6aeb8 100644 --- a/uv.lock +++ b/uv.lock @@ -815,6 +815,7 @@ dev = [ { name = "opentelemetry-sdk" }, { name = "packaging" }, { name = "pre-commit" }, + { name = "psutil" }, { name = "pydantic-ai-slim", extra = ["anthropic", "openai"], marker = "extra == 'extra-8-band-sdk-dev' or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-strands-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-pydantic-ai' and extra == 'extra-8-band-sdk-strands')" }, { name = "pydantic-settings" }, { name = "pyrefly" }, @@ -955,6 +956,9 @@ nooa = [ { name = "nooa", marker = "(python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant')" }, { name = "openai", version = "2.46.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant')" }, ] +omp = [ + { name = "psutil" }, +] opencode = [ { name = "httpx" }, { name = "mcp", version = "1.28.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-band-sdk-crewai' or extra == 'extra-8-band-sdk-dev-crewai' or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-strands-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-pydantic-ai' and extra == 'extra-8-band-sdk-strands')" }, @@ -1104,6 +1108,8 @@ requires-dist = [ { name = "pillow", marker = "extra == 'crewai'", specifier = ">=12.1.1" }, { name = "pillow", marker = "extra == 'dev-crewai'", specifier = ">=12.1.1" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.0.0" }, + { name = "psutil", marker = "extra == 'dev'", specifier = ">=6" }, + { name = "psutil", marker = "extra == 'omp'", specifier = ">=6" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pydantic-ai-slim", extras = ["anthropic", "openai"], marker = "extra == 'dev'", specifier = ">=2.55.0" }, { name = "pydantic-ai-slim", extras = ["anthropic", "openai"], marker = "extra == 'pydantic-ai'", specifier = ">=2.55.0" }, @@ -1191,7 +1197,7 @@ requires-dist = [ { name = "werkzeug", marker = "extra == 'dev-parlant'", specifier = ">=3.1.6" }, { name = "werkzeug", marker = "extra == 'parlant'", specifier = ">=3.1.6" }, ] -provides-extras = ["logging", "desktop", "codex", "opencode", "letta", "pydantic-ai", "anthropic", "langgraph", "claude-sdk", "copilot-sdk", "parlant", "crewai", "gemini", "a2a", "a2a-gateway", "a2a-gateway-demo", "acp", "slack", "bridge", "bridge-agentcore", "agentcore-runtime", "google-adk", "agno", "strands", "nooa", "dev", "dev-crewai", "dev-parlant", "dev-strands-nooa"] +provides-extras = ["logging", "desktop", "codex", "opencode", "letta", "pydantic-ai", "anthropic", "langgraph", "claude-sdk", "copilot-sdk", "parlant", "crewai", "gemini", "a2a", "a2a-gateway", "a2a-gateway-demo", "acp", "omp", "slack", "bridge", "bridge-agentcore", "agentcore-runtime", "google-adk", "agno", "strands", "nooa", "dev", "dev-crewai", "dev-parlant", "dev-strands-nooa"] [[package]] name = "band-sdk-core" @@ -6845,6 +6851,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + [[package]] name = "py-key-value-aio" version = "0.4.4"