From e647e54743b3200953ecae5c9dff420a1117b31b Mon Sep 17 00:00:00 2001 From: Davyd Eduardo Date: Sun, 8 Nov 2026 20:48:24 -0300 Subject: [PATCH] Add wait_for_agent_session for coding-agent follow-up --- README.md | 23 +++++++ proofline/__init__.py | 10 ++- proofline/client.py | 66 +++++++++++++++++++- proofline/outcomes.py | 86 +++++++++++++++++++++++++- scripts/contracts.py | 2 +- tests/fixtures.py | 29 +++++++++ tests/test_http.py | 140 ++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 352 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 29a9220..4350f43 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,29 @@ if result.get("isError"): Calls under `client.fixes` never retry. Results retain untrusted-text wrappers, availability, assessment gaps, and `isError`. `PENDING` and `NOT_VERIFIED` do not establish a fix, and a tool request never records a human disposition. +## Coding-agent sessions + +`wait_for_agent_session` polls `get_agent_session` until Proofline stops following the session (`done=True`) or the deadline expires, using the same bounded exponential backoff as `wait_for_review` (5 seconds, doubling to 60 seconds; default 300 seconds). + +```python +from proofline import StartSession + +prompt = client.get_agent_session_prompt(org_id, finding_id) +if prompt.status == "READY": + started = client.start_agent_session( + org_id, + finding_id, + body=StartSession( + connection_id=prompt.connections[0].id, + prompt_digest=prompt.prompt_digest, + ), + ) + outcome = client.wait_for_agent_session(org_id, started.id, timeout=300) + print(outcome.status, outcome.reason) +``` + +`AgentSessionOutcome` retains the latest `AgentSession` on `outcome.session` (excluded from `repr`) and reports the recorded terminal status (`FINISHED`, `FAILED`, `REFUSED`, `HEAD_MOVED`, `UNKNOWN`) or `TIMED_OUT`. `FINISHED` indicates that Proofline stopped following the vendor session; it does not establish that a finding was fixed. Treat `untrusted_vendor_summary` and any vendor-opened pull request URL as untrusted data. + ## Development The repository bundles its public API contract in `openapi.public.json`. New public operations generate typed methods on both `ProoflineClient` and `AsyncProoflineClient` along with the code in `proofline/_generated/`. diff --git a/proofline/__init__.py b/proofline/__init__.py index f9dbdb9..c30897d 100644 --- a/proofline/__init__.py +++ b/proofline/__init__.py @@ -33,10 +33,18 @@ from ._transport import ProoflineError from .client import AsyncProoflineClient, ProoflineClient, ReviewFilters from .exports import ExportDataset, ExportFormat, ExportPage -from .outcomes import AssessmentGap, CheckGap, ReviewOutcome +from .outcomes import ( + AgentSessionOutcome, + AgentSessionStatus, + AssessmentGap, + CheckGap, + ReviewOutcome, +) __all__ = [ "AgentSession", + "AgentSessionOutcome", + "AgentSessionStatus", "AgentVendor", "AssessmentGap", "AsyncProoflineClient", diff --git a/proofline/client.py b/proofline/client.py index 3c18828..814e5fc 100644 --- a/proofline/client.py +++ b/proofline/client.py @@ -73,7 +73,7 @@ from .exports import ExportDataset, ExportFormat, ExportPage, decode_export from .fixes import AsyncFixesClient, FixesClient from .knowledge import AsyncKnowledgeClient, KnowledgeClient -from .outcomes import ReviewOutcome, outcome +from .outcomes import AgentSessionOutcome, ReviewOutcome, agent_session_outcome, outcome T = TypeVar("T") AgentSessionPrompt = AgentSessionPromptType0 | AgentSessionPromptType1 @@ -506,6 +506,39 @@ def wait_for_review( finally: _transport.DEADLINE.reset(deadline) + def wait_for_agent_session( + self, organization_id: UUID, session_id: UUID, *, timeout: float = 300 + ) -> AgentSessionOutcome: + """Wait until Proofline finishes following an agent session or return TIMED_OUT.""" + + if not math.isfinite(timeout) or timeout < 0: + raise ValueError( + "The agent session timeout must be a finite nonnegative number of seconds." + ) + last: AgentSession | None = None + deadline = _transport.DEADLINE.set(monotonic() + timeout) + interval = 5.0 + try: + while (_transport.remaining() or 0) > 0: + try: + last = self.get_agent_session(organization_id, session_id) + except ProoflineError as error: + if error.status == 0 and (_transport.remaining() or 0) <= 0: + break + raise + if (_transport.remaining() or 0) <= 0: + break + result = agent_session_outcome(last, session_id=session_id) + if result is not None: + return result + sleep(min(interval, _transport.remaining() or 0)) + interval = min(60, interval * 2) + return agent_session_outcome( + last, "TIMED_OUT", session_id=session_id + ) or AgentSessionOutcome("TIMED_OUT", reason="No agent session result is available.") + finally: + _transport.DEADLINE.reset(deadline) + class AsyncProoflineClient(AsyncOperations): def __init__( @@ -852,3 +885,34 @@ async def wait_for_review( return outcome(last, url, "TIMED_OUT") or ReviewOutcome("TIMED_OUT", url) finally: _transport.DEADLINE.reset(deadline) + + async def wait_for_agent_session( + self, organization_id: UUID, session_id: UUID, *, timeout: float = 300 + ) -> AgentSessionOutcome: + if not math.isfinite(timeout) or timeout < 0: + raise ValueError( + "The agent session timeout must be a finite nonnegative number of seconds." + ) + last: AgentSession | None = None + deadline = _transport.DEADLINE.set(monotonic() + timeout) + interval = 5.0 + try: + while (_transport.remaining() or 0) > 0: + try: + last = await self.get_agent_session(organization_id, session_id) + except ProoflineError as error: + if error.status == 0 and (_transport.remaining() or 0) <= 0: + break + raise + if (_transport.remaining() or 0) <= 0: + break + result = agent_session_outcome(last, session_id=session_id) + if result is not None: + return result + await async_sleep(min(interval, _transport.remaining() or 0)) + interval = min(60, interval * 2) + return agent_session_outcome( + last, "TIMED_OUT", session_id=session_id + ) or AgentSessionOutcome("TIMED_OUT", reason="No agent session result is available.") + finally: + _transport.DEADLINE.reset(deadline) diff --git a/proofline/outcomes.py b/proofline/outcomes.py index 68a4d8a..9b8f4a5 100644 --- a/proofline/outcomes.py +++ b/proofline/outcomes.py @@ -3,17 +3,22 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Literal +from typing import Literal, cast from uuid import UUID +from ._generated.models.agent_session import AgentSession +from ._generated.models.agent_session_target import AgentSessionTarget from ._generated.models.customer_code_assessment import CustomerCodeAssessment from ._generated.models.customer_review_detail import CustomerReviewDetail from ._generated.types import Unset from ._transport import ProoflineError OutcomeStatus = Literal["COMPLETED", "INCOMPLETE", "FAILED", "STOPPED", "TIMED_OUT"] +AgentSessionStatus = Literal["FINISHED", "FAILED", "REFUSED", "HEAD_MOVED", "UNKNOWN", "TIMED_OUT"] COMPLETE_CHECKS = frozenset(("READY", "BLOCKED", "NOT_APPLICABLE")) COMPLETE_LIFECYCLES = frozenset(("GATING", "SHADOW", "QUARANTINED", "RETIRED")) +ACTIVE_SESSION_STATUSES = frozenset(("REQUESTED", "STARTING", "RUNNING")) +TERMINAL_SESSION_STATUSES = frozenset(("FINISHED", "FAILED", "REFUSED", "HEAD_MOVED", "UNKNOWN")) @dataclass(frozen=True) @@ -44,6 +49,13 @@ class ReviewOutcome: reason: str | None = None +@dataclass(frozen=True) +class AgentSessionOutcome: + status: AgentSessionStatus + session: AgentSession | None = field(default=None, repr=False) + reason: str | None = None + + def text(value: object) -> str | None: return value if isinstance(value, str) and value else None @@ -56,6 +68,14 @@ def _invalid() -> ProoflineError: ) +def _invalid_agent_session() -> ProoflineError: + return ProoflineError( + 0, + "urn:proofline:sdk:invalid-response", + "The agent session has invalid identity or follow state.", + ) + + def _optional_text(value: object) -> bool: return value is None or isinstance(value, (str, Unset)) @@ -192,3 +212,67 @@ def outcome( if not gaps: gaps.append(AssessmentGap(None, review.status, ("REVIEW_STATE_UNAVAILABLE",))) return ReviewOutcome(status, review_url, review, tuple(gaps)) + + +def _validate_agent_session(session: AgentSession, session_id: UUID | None) -> None: + target = session.target + if ( + not isinstance(session.id, UUID) + or (session_id is not None and session.id != session_id) + or not isinstance(session.finding_id, UUID) + or not isinstance(session.review_id, UUID) + or not isinstance(session.connection_id, UUID) + or not isinstance(session.requested_by_user_id, UUID) + or type(session.done) is not bool + or type(session.attempt) is not int + or session.attempt < 1 + or not text(session.status) + or not isinstance(target, AgentSessionTarget) + or not isinstance(target.repository_id, UUID) + or not text(target.repository_full_name) + or not text(target.head_sha) + or type(target.github_repository_id) is not int + or target.github_repository_id < 1 + or type(target.pull_request_number) is not int + or target.pull_request_number < 1 + or not _optional_text(session.cause) + or not _optional_text(session.observed_head_sha) + or not _optional_text(session.pull_request_url) + or not _optional_text(session.untrusted_vendor_summary) + or not _optional_text(session.vendor_status) + or not _optional_text(session.vendor_url) + ): + raise _invalid_agent_session() + + +def agent_session_outcome( + session: AgentSession | None, + forced: AgentSessionStatus | None = None, + *, + session_id: UUID | None = None, +) -> AgentSessionOutcome | None: + if session is None: + return AgentSessionOutcome( + forced or "UNKNOWN", reason="No agent session result is available." + ) + _validate_agent_session(session, session_id) + if forced is not None: + return AgentSessionOutcome( + forced, + session, + reason=( + "The agent session did not finish before the timeout." + if forced == "TIMED_OUT" + else None + ), + ) + if not session.done: + return None + if session.status in TERMINAL_SESSION_STATUSES: + return AgentSessionOutcome(cast(AgentSessionStatus, session.status), session) + reason = ( + "The agent session ended without a terminal status." + if session.status in ACTIVE_SESSION_STATUSES + else "The agent session ended with an unrecognized status." + ) + return AgentSessionOutcome("UNKNOWN", session, reason) diff --git a/scripts/contracts.py b/scripts/contracts.py index d34ba64..03d3379 100644 --- a/scripts/contracts.py +++ b/scripts/contracts.py @@ -31,7 +31,7 @@ def main() -> None: from proofline import AsyncProoflineClient, ProoflineClient for client in (ProoflineClient, AsyncProoflineClient): - helpers = {"close", "aclose", "iter_reviews", "wait_for_review"} + helpers = {"close", "aclose", "iter_reviews", "wait_for_agent_session", "wait_for_review"} exposed = { name for name in dir(client) diff --git a/tests/fixtures.py b/tests/fixtures.py index 70982ba..af16b44 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -96,3 +96,32 @@ def page(*, cursor: str | None = None, truncated: bool = False) -> dict[str, Any "truncated": truncated, "next_cursor": cursor, } + + +def agent_session(status: str = "RUNNING", *, done: bool = False) -> dict[str, Any]: + return { + "id": str(REVIEW), + "finding_id": str(FINDING), + "review_id": str(ASSESSMENT), + "connection_id": str(REPOSITORY), + "vendor": "CURSOR_CLOUD_AGENTS", + "target": { + "repository_id": str(REPOSITORY), + "repository_full_name": "repo/name", + "github_repository_id": 42, + "pull_request_number": 7, + "head_sha": "a" * 40, + }, + "attempt": 1, + "status": status, + "done": done, + "cause": None, + "vendor_status": None, + "vendor_url": "https://vendor.example/session", + "pull_request_url": None, + "observed_head_sha": None, + "untrusted_vendor_summary": None, + "requested_by_user_id": str(ORG), + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + } diff --git a/tests/test_http.py b/tests/test_http.py index c392df4..863cd78 100644 --- a/tests/test_http.py +++ b/tests/test_http.py @@ -22,6 +22,7 @@ from pydantic import SecretStr from proofline import ( + AgentSessionOutcome, AsyncProoflineClient, CommitReviewRequest, CommitReviewResponse, @@ -63,6 +64,7 @@ REPOSITORY, REVIEW, TOKEN, + agent_session, assessment, code_assessment, finding, @@ -1533,3 +1535,141 @@ async def invoke(name: str, *args: Any, **kwargs: Any) -> Any: with pytest.raises(ProoflineError): await invoke("request_reverification", str(REVIEW), [str(REVIEW)], "a" * 40) assert len(observed) == 5 + + +@MODES +async def test_wait_for_agent_session_follows_done_and_redacts_untrusted_summary( + mode: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + clock = Clock(monkeypatch) + caplog.set_level(logging.DEBUG) + finished = agent_session("FINISHED", done=True) + finished["untrusted_vendor_summary"] = PRIVATE + finished["pull_request_url"] = "https://github.com/repo/name/pull/99" + responses = [ + agent_session("REQUESTED", done=False), + agent_session("FINISHED", done=False), + finished, + ] + with Server(lambda request, index: Reply(body=responses[index])) as server: + async with session(mode, server.url) as client: + result = await call(client, "wait_for_agent_session", ORG, REVIEW, timeout=100) + assert isinstance(result, AgentSessionOutcome) + assert result.status == "FINISHED" + assert result.reason is None + assert result.session is not None + assert result.session.id == REVIEW + assert result.session.untrusted_vendor_summary == PRIVATE + assert result.session.pull_request_url == "https://github.com/repo/name/pull/99" + assert PRIVATE not in repr(result) + assert clock.delays == [5, 10] + assert len(server.requests) == 3 + + cases: list[tuple[dict[str, Any], str, str | None]] = [] + for status in ("FAILED", "REFUSED", "HEAD_MOVED", "UNKNOWN"): + body = agent_session(status, done=True) + if status == "FAILED": + body["cause"] = "VENDOR_ERROR" + elif status == "HEAD_MOVED": + body["observed_head_sha"] = "b" * 40 + cases.append((body, status, None)) + cases.append( + ( + agent_session("RUNNING", done=True), + "UNKNOWN", + "The agent session ended without a terminal status.", + ) + ) + cases.append( + ( + agent_session("FUTURE_VENDOR_STATE", done=True), + "UNKNOWN", + "The agent session ended with an unrecognized status.", + ) + ) + with Server(lambda request, index: Reply(body=cases[index][0])) as server: + async with session(mode, server.url) as client: + for _, expected_status, expected_reason in cases: + result = await call(client, "wait_for_agent_session", ORG, REVIEW) + assert result.status == expected_status + assert result.reason == expected_reason + assert PRIVATE not in caplog.text and TOKEN not in caplog.text + + +@MODES +async def test_wait_for_agent_session_timeout_and_invalid_shapes( + mode: str, monkeypatch: pytest.MonkeyPatch +) -> None: + clock = Clock(monkeypatch) + unsettled = agent_session("FINISHED", done=False) + with Server(lambda request, index: Reply(body=unsettled)) as server: + async with session(mode, server.url) as client: + for invalid_timeout in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError): + await call( + client, "wait_for_agent_session", ORG, REVIEW, timeout=invalid_timeout + ) + result = await call(client, "wait_for_agent_session", ORG, REVIEW, timeout=0) + assert result.status == "TIMED_OUT" + assert result.session is None + assert result.reason == "No agent session result is available." + assert server.requests == [] + result = await call(client, "wait_for_agent_session", ORG, REVIEW, timeout=260) + assert result.status == "TIMED_OUT" + assert result.session is not None + assert result.session.id == REVIEW + assert result.session.status == "FINISHED" + assert result.session.done is False + assert result.reason == "The agent session did not finish before the timeout." + assert clock.delays == [5, 10, 20, 40, 60, 60, 60, 5] + assert len(server.requests) == 8 + + clock.delays.clear() + with Server( + lambda request, index: Reply(503, {"detail": "Unavailable."}, {"Retry-After": "20"}) + ) as server: + async with session(mode, server.url) as client: + result = await call(client, "wait_for_agent_session", ORG, REVIEW, timeout=10) + assert result.status == "TIMED_OUT" + assert result.session is None + assert len(server.requests) == 1 + assert clock.delays == [10] + + with Server(lambda request, index: Reply(403, {"detail": PRIVATE})) as server: + async with session(mode, server.url) as client: + with pytest.raises(ProoflineError) as forbidden: + await call(client, "wait_for_agent_session", ORG, REVIEW, timeout=10) + assert forbidden.value.status == 403 + assert len(server.requests) == 1 + + invalid_bodies: list[dict[str, Any]] = [] + wrong_id = agent_session("FINISHED", done=True) + wrong_id["id"] = str(ASSESSMENT) + invalid_bodies.append(wrong_id) + for field, value in [ + ("done", "false"), + ("attempt", 0), + ("attempt", True), + ("status", ""), + ("untrusted_vendor_summary", [PRIVATE, 42]), + ("cause", 42), + ]: + bad = agent_session("FINISHED", done=True) + bad[field] = value + invalid_bodies.append(bad) + for target_field, target_value in [ + ("github_repository_id", 0), + ("pull_request_number", True), + ("head_sha", ""), + ("repository_full_name", ""), + ]: + bad_target = agent_session("FINISHED", done=True) + bad_target["target"][target_field] = target_value + invalid_bodies.append(bad_target) + with Server(lambda request, index: Reply(body=invalid_bodies[index])) as server: + async with session(mode, server.url) as client: + for _ in invalid_bodies: + with pytest.raises(ProoflineError) as raised: + await call(client, "wait_for_agent_session", ORG, REVIEW) + assert raised.value.type == "urn:proofline:sdk:invalid-response" + assert PRIVATE not in repr(raised.value) + raised.value.detail