diff --git a/backend/druks/accounts/dependencies.py b/backend/druks/accounts/dependencies.py index 232b5321..cf735bb8 100644 --- a/backend/druks/accounts/dependencies.py +++ b/backend/druks/accounts/dependencies.py @@ -20,11 +20,11 @@ _bearer_scheme = HTTPBearer(auto_error=False, scheme_name="personalAccessToken") -def resolve_pat_account(credentials: HTTPAuthorizationCredentials | None) -> Account: +async def resolve_pat_account(credentials: HTTPAuthorizationCredentials | None) -> Account: """A present Authorization must authenticate — never a fall-through.""" if credentials: try: - return PersonalAccessToken.authenticate(credentials.credentials).account + return (await PersonalAccessToken.authenticate(credentials.credentials)).account except InvalidPatError as error: raise HTTPException( status_code=401, @@ -38,10 +38,10 @@ def resolve_pat_account(credentials: HTTPAuthorizationCredentials | None) -> Acc ) -def resolve_single_operator() -> Account | None: +async def resolve_single_operator() -> Account | None: """None while zero accounts exist (setup); more than one refuses rather than guesses.""" - operators = Account.list_non_system() + operators = await Account.list_non_system() if len(operators) > 1: raise AuthConfigurationError( f"auth mode 'none' expects exactly one operator account, found " @@ -56,16 +56,16 @@ async def _resolve_operator(connection: HTTPConnection) -> Account | None: connection, not a request, so a WebSocket upgrade resolves the same way.""" settings = connection.app.state.settings if settings.identity.mode == "none": - return resolve_single_operator() + return await resolve_single_operator() values = connection.headers.getlist(settings.identity.header) if len(values) == 1 and (asserted := values[0].strip()): if settings.identity.mode == "header": - return Account.get_or_create(asserted) + return await Account.get_or_create(asserted) try: email = await verify_assertion(asserted, settings) except InvalidAssertionError as error: raise HTTPException(status_code=401, detail=str(error)) from error - return Account.get_or_create(email) + return await Account.get_or_create(email) raise HTTPException( status_code=401, detail=f"The edge must assert exactly one nonblank {settings.identity.header} identity.", @@ -102,7 +102,7 @@ async def current_account( """The Bearer PAT when Authorization is present — present-but-empty still challenges — else the session identity.""" if "Authorization" in request.headers: - account = resolve_pat_account(bearer) + account = await resolve_pat_account(bearer) else: account = await _resolve_operator(request) if not account: @@ -154,7 +154,7 @@ async def current_account_or_setup( """PAT-first identity that reads none/zero setup as None instead of refusing — ``/api/auth/me`` only.""" if "Authorization" in request.headers: - account = resolve_pat_account(bearer) + account = await resolve_pat_account(bearer) else: account = await _resolve_operator(request) token = current_account_id.set(account.id if account else None) diff --git a/backend/druks/accounts/models.py b/backend/druks/accounts/models.py index 8ac1f094..1ef29418 100644 --- a/backend/druks/accounts/models.py +++ b/backend/druks/accounts/models.py @@ -38,35 +38,35 @@ class Account(Base, Uuid7Pk): created_at: Mapped[datetime] = mapped_column(default=Base.utc_now) @classmethod - def get(cls, account_id: str, *, exclude_system: bool = False) -> "Account | None": + async def get(cls, account_id: str, *, exclude_system: bool = False) -> "Account | None": if exclude_system and account_id == SYSTEM_ACCOUNT_ID: return - return db_session().get(cls, account_id) + return await db_session().get(cls, account_id) @classmethod - def get_for_username(cls, username: str) -> "Account | None": - return db_session().scalar(select(cls).where(cls.username == username)) + async def get_for_username(cls, username: str) -> "Account | None": + return await db_session().scalar(select(cls).where(cls.username == username)) @classmethod - def get_or_create(cls, username: str) -> "Account": + async def get_or_create(cls, username: str) -> "Account": """Concurrency-safe lookup-or-create: racing requests both INSERT with ON CONFLICT DO NOTHING, then converge on the one row through the canonical CITEXT lookup.""" - account = cls.get_for_username(username) + account = await cls.get_for_username(username) if account: return account session = db_session() - session.execute( + await session.execute( insert(cls) .values(username=username) .on_conflict_do_nothing(index_elements=["username"]) ) - return session.scalars(select(cls).where(cls.username == username)).one() + return (await session.scalars(select(cls).where(cls.username == username))).one() @classmethod - def list_non_system(cls) -> list["Account"]: + async def list_non_system(cls) -> list["Account"]: stmt = select(cls).where(cls.username != SYSTEM_ACCOUNT_ID).order_by(cls.created_at) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) def _hash_token(token: str) -> bytes: @@ -108,24 +108,24 @@ def status(self) -> str: return "active" @classmethod - def get(cls, pat_id: str) -> "PersonalAccessToken | None": - return db_session().get(cls, pat_id) + async def get(cls, pat_id: str) -> "PersonalAccessToken | None": + return await db_session().get(cls, pat_id) @classmethod - def get_for_prefix(cls, prefix: str) -> "PersonalAccessToken | None": - return db_session().scalar(select(cls).where(cls.token_prefix == prefix)) + async def get_for_prefix(cls, prefix: str) -> "PersonalAccessToken | None": + return await db_session().scalar(select(cls).where(cls.token_prefix == prefix)) @classmethod - def list_for_account(cls, account_id: str) -> list["PersonalAccessToken"]: + async def list_for_account(cls, account_id: str) -> list["PersonalAccessToken"]: stmt = select(cls).where(cls.account_id == account_id).order_by(cls.created_at.desc()) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def create(cls, *, account_id: str, name: str) -> "tuple[PersonalAccessToken, str]": + async def create(cls, *, account_id: str, name: str) -> "tuple[PersonalAccessToken, str]": """Mint ``account_id`` a token; returns (row, plaintext). The plaintext is shown exactly once — only its hash lands in the row.""" prefix = _new_prefix() - while cls.get_for_prefix(prefix): + while await cls.get_for_prefix(prefix): prefix = _new_prefix() secret = base64.urlsafe_b64encode(secrets.token_bytes(PAT_SECRET_BYTES)) token = f"{PAT_TOKEN_TAG}_{prefix}_{secret.rstrip(b'=').decode()}" @@ -141,16 +141,16 @@ def create(cls, *, account_id: str, name: str) -> "tuple[PersonalAccessToken, st ) session = db_session() session.add(row) - session.flush() + await session.flush() return row, token @classmethod - def authenticate(cls, credential: str) -> "PersonalAccessToken": + async def authenticate(cls, credential: str) -> "PersonalAccessToken": """Resolve a presented bearer credential to its live row — the one authentication door for both HTTP and MCP — or raise InvalidPatError. Stamps last_used_at, at most hourly.""" prefix, _, _ = credential.removeprefix(f"{PAT_TOKEN_TAG}_").partition("_") - row = cls.get_for_prefix(prefix) + row = await cls.get_for_prefix(prefix) if not row: raise InvalidPatError("Not a recognized personal access token.") if not hmac.compare_digest(_hash_token(credential), row.token_hash): @@ -162,10 +162,10 @@ def authenticate(cls, credential: str) -> "PersonalAccessToken": now = Base.utc_now() if not row.last_used_at or now - row.last_used_at >= PAT_LAST_USED_RESOLUTION: row.last_used_at = now - db_session().flush() + await db_session().flush() return row - def revoke(self) -> None: + async def revoke(self) -> None: # Keep the first revocation instant — a repeat revoke changes nothing. self.revoked_at = self.revoked_at or Base.utc_now() - db_session().flush() + await db_session().flush() diff --git a/backend/druks/accounts/routes.py b/backend/druks/accounts/routes.py index 1456eb7e..5b135061 100644 --- a/backend/druks/accounts/routes.py +++ b/backend/druks/accounts/routes.py @@ -18,7 +18,7 @@ async def get_identity( account=AccountResponse.model_validate(account) if account else None, # An account needs onboarding exactly while it has no harness # connection; none/zero is onboarding before the account exists. - onboarding_required=not (account and HarnessConnection.list_for_account(account.id)), + onboarding_required=not (account and await HarnessConnection.list_for_account(account.id)), ) @@ -26,7 +26,7 @@ async def get_identity( async def list_pats( account: Account = Depends(current_session_account), ) -> list[PersonalAccessToken]: - return PersonalAccessToken.list_for_account(account.id) + return await PersonalAccessToken.list_for_account(account.id) @router.post("/personal-tokens") @@ -38,7 +38,7 @@ async def create_pat( if name and len(name) <= PAT_NAME_LENGTH: # The plaintext, handed back exactly once — only its hash is stored, # and the new row surfaces through the list. - _, token = PersonalAccessToken.create(account_id=account.id, name=name) + _, token = await PersonalAccessToken.create(account_id=account.id, name=name) return {"token": token} raise HTTPException( status_code=422, @@ -52,9 +52,9 @@ async def create_pat( async def revoke_pat( pat_id: str, account: Account = Depends(current_session_account) ) -> PersonalAccessToken: - pat = PersonalAccessToken.get(pat_id) + pat = await PersonalAccessToken.get(pat_id) if pat and pat.account_id == account.id: - pat.revoke() + await pat.revoke() return pat # One shape for missing and foreign — existence stays account-scoped. raise HTTPException(status_code=404, detail="No such token.") diff --git a/backend/druks/agents.py b/backend/druks/agents.py index 26ebf5fd..7dbab869 100644 --- a/backend/druks/agents.py +++ b/backend/druks/agents.py @@ -123,16 +123,16 @@ def __set_name__(self, owner: type, attr: str) -> None: # override → the agent's declared value → the operator's global default. # ``run`` uses these; callers that drive the harness themselves call them # directly. - def get_model_name(self) -> str: - return SettingsOverride.agent_model(self.id, self.model).value + async def get_model_name(self) -> str: + return (await SettingsOverride.agent_model(self.id, self.model)).value - def get_effort(self) -> str: - harness = get_harness_for_model(self.get_model_name()).name - return SettingsOverride.agent_effort(self.id, self.effort, harness).value + async def get_effort(self) -> str: + harness = (await get_harness_for_model(await self.get_model_name())).name + return (await SettingsOverride.agent_effort(self.id, self.effort, harness)).value - def get_timeout(self) -> int: - harness = get_harness_for_model(self.get_model_name()).name - resolved = SettingsOverride.agent_timeout(self.id, self.timeout, harness).value + async def get_timeout(self) -> int: + harness = (await get_harness_for_model(await self.get_model_name())).name + resolved = (await SettingsOverride.agent_timeout(self.id, self.timeout, harness)).value # Capped so a single call always fits inside a fresh sandbox lease. return min(resolved, MAX_AGENT_TIMEOUT_SECONDS) @@ -187,11 +187,11 @@ async def _quota_retry_wait() -> float: # Runs as its own step: the body does no IO, and replay reuses the # recorded wait instead of re-reading the scrape. async with step_session(): - harness = get_harness_for_model(self.get_model_name()) + harness = await get_harness_for_model(await self.get_model_name()) # The scrape belongs to the charged connection — its account # differs from the run's on fallback. - connection = HarnessConnection.lookup(harness.name, workflow.account_id) - scrape = UsageScrape.latest_for(harness.name, connection.account_id) + connection = await HarnessConnection.lookup(harness.name, workflow.account_id) + scrape = await UsageScrape.latest_for(harness.name, connection.account_id) if scrape: now = datetime.now(UTC) reset = scrape.soonest_reset_after(now) @@ -253,18 +253,18 @@ async def _run(self, *, workflow_id: str, **context: Any) -> Any: the harness. ``__call__`` handles the durable wrapping + nesting.""" if not self.prompt: raise WorkflowError(f"agent {self.id!r} has no prompt template to render") - model = self.get_model_name() - harness = get_harness_for_model(model) + model = await self.get_model_name() + harness = await get_harness_for_model(model) workflow = current_workflow.get() # Refusing an unservable call here beats provisioning a VM and # 401ing mid-run. - connection = HarnessConnection.lookup(harness.name, workflow.account_id) + connection = await HarnessConnection.lookup(harness.name, workflow.account_id) # Plain snapshots: the commits below expire the ORM row mid-flight. connection_id, charged_account_id = connection.id, connection.account_id # An agent call is a durability boundary — its effects don't roll back — # so commit here rather than hold the step's connection idle through the # minutes of provisioning and the run. - db_session().commit() + await db_session().commit() settings = load_settings() artifact_dir = settings.artifacts_dir / f"run-{workflow_id}" @@ -290,7 +290,7 @@ async def _run(self, *, workflow_id: str, **context: Any) -> Any: prompt_context.setdefault("workspace", runner) prompt = await render_prompt(self.prompt, **prompt_context) await set_run_phase("agent_running") - AgentCall.start( + await AgentCall.start( engine, call_id=call_id, run_id=workflow_id, @@ -310,16 +310,16 @@ async def _run(self, *, workflow_id: str, **context: Any) -> Any: account_id=workflow.account_id, ) except BaseException as error: - AgentCall.fail(engine, call_id=call_id, error=error) + await AgentCall.fail(engine, call_id=call_id, error=error) raise - AgentCall.finish(engine, call_id=call_id, result=result) + await AgentCall.finish(engine, call_id=call_id, result=result) if result.error: raise result.error output = self.contract.model_validate(result.output) if spec := output.get_artifact(): - Artifact.record(call_dir=artifact_dir / call_id, call_id=call_id, **spec) + await Artifact.record(call_dir=artifact_dir / call_id, call_id=call_id, **spec) return output.to_result() async def _execute( @@ -340,8 +340,8 @@ async def _execute( prompt=prompt, schema=schema, agent=self.id, - effort=self.get_effort(), - timeout=self.get_timeout(), + effort=await self.get_effort(), + timeout=await self.get_timeout(), artifact_dir=artifact_dir, call_id=call_id, include_plugins=self.include_plugins, diff --git a/backend/druks/api/artifacts.py b/backend/druks/api/artifacts.py index 00b3fd1a..367e05b5 100644 --- a/backend/druks/api/artifacts.py +++ b/backend/druks/api/artifacts.py @@ -11,10 +11,10 @@ async def get_artifact(artifact_id: str) -> ArtifactContent: # A call's renderable output, reached through the call that produced it — the in-app # review fetches this to render the plan beside its controls. - artifact = db_session().get(Artifact, artifact_id) + artifact = await db_session().get(Artifact, artifact_id) if not artifact: raise HTTPException(status.HTTP_404_NOT_FOUND, "artifact not found") - call = AgentCall.get(artifact.agent_call_id) + call = await AgentCall.get(artifact.agent_call_id) path = call.get_file_path(artifact.path) if not path: raise HTTPException(status.HTTP_404_NOT_FOUND, "artifact content missing") diff --git a/backend/druks/api/health_status.py b/backend/druks/api/health_status.py index a907adf9..2808deb6 100644 --- a/backend/druks/api/health_status.py +++ b/backend/druks/api/health_status.py @@ -7,9 +7,9 @@ from druks.webhooks.deliveries import last_delivery_at -def _spend_for_local_today(*, timezone_name: str, now: datetime) -> tuple[float, int]: +async def _spend_for_local_today(*, timezone_name: str, now: datetime) -> tuple[float, int]: _, local_start = operator_local_day(timezone_name, now) - return AgentCall.total_run_spend_between( + return await AgentCall.total_run_spend_between( start=local_start.astimezone(UTC), end=(local_start + timedelta(days=1)).astimezone(UTC), ) @@ -17,7 +17,9 @@ def _spend_for_local_today(*, timezone_name: str, now: datetime) -> tuple[float, async def build_health() -> api_schemas.DashboardHealth: now = datetime.now(UTC) - spend, tokens = _spend_for_local_today(timezone_name=UserSettings.get().timezone, now=now) + spend, tokens = await _spend_for_local_today( + timezone_name=(await UserSettings.get()).timezone, now=now + ) # GitHub is the code host and is always present. sources = ["github"] return api_schemas.DashboardHealth( diff --git a/backend/druks/api/runs.py b/backend/druks/api/runs.py index cdf685f1..c4b4ef22 100644 --- a/backend/druks/api/runs.py +++ b/backend/druks/api/runs.py @@ -22,7 +22,7 @@ async def resume_run(run_id: Annotated[str, Path(alias="run")], body: ResumeRequest) -> None: # The in-app half of a gate: the operator answers the parked run from Druks # (external gates resume through their own webhook). - run = Run.get(run_id) + run = await Run.get(run_id) if not run: raise HTTPException(status.HTTP_404_NOT_FOUND, "run not found") ask = run.input_request @@ -65,7 +65,7 @@ async def cancel_run( ) -> CancelRunResponse: """Cancel an active run, recording the reason as its failure; a repeat cancel reports already_cancelled.""" - run = Run.get(run_id) + run = await Run.get(run_id) if not run: raise RunNotFound(run_id) if run.state == RunState.CANCELLED.value: @@ -94,15 +94,15 @@ async def retry_run( ) -> RetryRunResponse: """Rerun a failed run from the step that killed it, reusing every completed step.""" - run = Run.get(run_id) + run = await Run.get(run_id) if not run: raise RunNotFound(run_id) if run.state != RunState.FAILED.value: raise RunNotFailed(run_id) - subject = run.subject + subject = await run.get_subject() if subject: - latest = Run.get_latest_for_subject(subject["type"], subject["id"]) + latest = await Run.get_latest_for_subject(subject["type"], subject["id"]) if latest and latest.is_active: raise SubjectBusy(latest.id) diff --git a/backend/druks/api/server.py b/backend/druks/api/server.py index 1d10629d..c16901f0 100644 --- a/backend/druks/api/server.py +++ b/backend/druks/api/server.py @@ -23,7 +23,12 @@ from druks.browser.exceptions import BrowserApiError from druks.browser.routes import router as browser_sessions_router from druks.core.templates import render_page -from druks.database import configure_session, create_engine_from_url, db_session, session_scope +from druks.database import ( + configure_session, + create_async_engine_from_url, + db_session, + session_scope, +) from druks.durable.engine import init_dbos, launch, shutdown from druks.durable.exceptions import AgentCallNotFound from druks.events.routes import router as events_router @@ -49,7 +54,7 @@ def configure_state(app: FastAPI, settings: Settings) -> None: ensure_data_dirs(settings) app.state.settings = settings - app.state.engine = create_engine_from_url(settings.database_url) + app.state.engine = create_async_engine_from_url(settings.database_url) # Bind the ambient (``scoped_session``) factory to this engine so # request handlers can use ``db_session()`` without per-call setup. configure_session(app.state.engine) @@ -57,12 +62,14 @@ def configure_state(app: FastAPI, settings: Settings) -> None: @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: - # Tests pre-populate ``app.state.settings`` before lifespan fires (or - # skip lifespan entirely by constructing ``TestClient(app)`` without - # ``with``). Production hits this branch and reads env config. + # Tests pre-populate ``app.state.settings`` before lifespan fires (their + # engine is a fixture-owned connection this lifespan must not dispose). + # Production hits this branch and reads env config. + created_engine = None if not hasattr(app.state, "settings"): settings = load_settings() configure_state(app, settings) + created_engine = app.state.engine # uvicorn runs this module directly and never calls setup_logging, so # app loggers default to WARNING-only and INFO is dropped; the web # process configures it here. @@ -75,13 +82,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # A drifted none-mode install (more than one operator account) must # refuse at boot, not per request; the per-request resolver repeats # the check for drift that happens while running. - with session_scope(app.state.engine): - resolve_single_operator() + async with session_scope(app.state.engine): + await resolve_single_operator() # DBOS runs embedded here: this process both serves HTTP and executes # durable workflows. Tests pre-populate app.state.settings and never # reach here — they drive DBOS through their own fixtures. init_dbos() - launch() + await launch() # Each app converges its own runtime state (e.g. schedules) here, after # DBOS is live. A failing hook is logged, not fatal — one app can't wedge boot. for registered_app in iter_apps(): @@ -96,33 +103,36 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: yield finally: shutdown() - engine = getattr(app.state, "engine", None) - if engine: - engine.dispose() + if created_engine: + await created_engine.dispose() await close_client() async def _release_db_session() -> AsyncIterator[None]: - """Commit the request's scoped DB session on success, roll back on error, - then release it — one transaction per request. Model writes ``flush()`` - without committing, so this is the commit boundary. FastAPI runs this - yield-dependency's teardown in the *same* asyncio task as the endpoint - (the scoped_session is keyed by that task), so it acts on exactly the - session this request opened. Frontend responses (the SPA, an app's - dist/) run app dependencies too, so only touch the registry when the - request actually opened a session — never open one just to commit nothing. - """ + """Bind a fresh DB session for the request and commit it on success, roll + back on error — one transaction per request. Model writes ``flush()`` + without committing, so this is the commit boundary. The session is the + request's own, never an ambient one already on this task (the test client + runs requests on the caller's task), and the prior binding is restored + after. The session object is lazy — a frontend response (the SPA, an + app's dist/) that never touches the DB opens no connection, and the + commit is a no-op.""" + previous = db_session() if db_session.registry.has() else None + session = db_session.session_factory() + db_session.registry.set(session) try: yield except BaseException: - if db_session.registry.has(): - db_session().rollback() + await session.rollback() raise else: - if db_session.registry.has(): - db_session().commit() + await session.commit() finally: - db_session.remove() + await session.close() + if previous is not None: + db_session.registry.set(previous) + else: + db_session.registry.clear() def _mcp_lifespan(app: FastAPI) -> AbstractAsyncContextManager[Mapping[str, Any] | None]: diff --git a/backend/druks/api/subjects.py b/backend/druks/api/subjects.py index b90bb2ea..58de11e0 100644 --- a/backend/druks/api/subjects.py +++ b/backend/druks/api/subjects.py @@ -21,7 +21,7 @@ async def list_open_subjects() -> OpenSubjectsResponse: """Every subject with open work — each nesting its open workflows, the newest run of one kind that is scheduled, running, parked, or failed. At most 50 workflows.""" - rows = db_session().execute(Run.get_open_subjects().limit(_WORKFLOW_ROWS)).all() + rows = (await db_session().execute(Run.get_open_subjects().limit(_WORKFLOW_ROWS))).all() grouped = {} for row in rows: grouped.setdefault(row.subject_type, {}).setdefault(row.subject_id, []).append(row) diff --git a/backend/druks/apps/base.py b/backend/druks/apps/base.py index 59db624f..94f2efd9 100644 --- a/backend/druks/apps/base.py +++ b/backend/druks/apps/base.py @@ -1,6 +1,6 @@ import importlib.util import re -from collections.abc import Callable +from collections.abc import Callable, Coroutine from pathlib import Path from types import ModuleType from typing import TYPE_CHECKING, Annotated, Any, ClassVar @@ -32,8 +32,8 @@ from druks.workflows import Workflow # A check the app owns returns a verdict on one of its own preconditions - # using the same ``CheckResult`` shape as a core check. - Check = Callable[[], CheckResult] + # using the same ``CheckResult`` shape as a core check; sync or async. + Check = Callable[[], "CheckResult | Coroutine[Any, Any, CheckResult]"] # An app name keys the ``/api/`` namespace, the ``alembic_version_`` # table, the ``_`` table prefix, and ``app::`` settings — so it must @@ -135,14 +135,14 @@ def __init_subclass__(cls, **kwargs: Any) -> None: cls.settings_model = declared @classmethod - def settings(cls) -> AppSettings: + async def settings(cls) -> AppSettings: """The app's settings, resolved through the override store keyed by app name. Raises if the app declares no ``Settings``.""" model = cls.settings_model if not model: raise TypeError(f"app {cls.name!r} declares no Settings") values = { - name: SettingsOverride.app_setting( + name: await SettingsOverride.app_setting( cls.name, name, field.default, @@ -153,7 +153,7 @@ def settings(cls) -> AppSettings: return model.model_validate(values) @classmethod - def override_setting(cls, field: str, value: Any) -> None: + async def override_setting(cls, field: str, value: Any) -> None: """An operator's override for one declared setting; ``None`` clears it back to the declared default. Raises ``ValueError`` so the API layer can 422 it.""" model = cls.settings_model @@ -161,8 +161,8 @@ def override_setting(cls, field: str, value: Any) -> None: raise ValueError(f"Unknown {cls.name} setting {field!r}") if value is not None: value = coerce_setting_value(model, field, value) - validate_setting_override(model, cls.settings().model_dump(), field, value) - SettingsOverride.set_app_setting( + validate_setting_override(model, (await cls.settings()).model_dump(), field, value) + await SettingsOverride.set_app_setting( cls.name, field, value, @@ -345,7 +345,7 @@ async def get_transcript( raise HTTPException( status.HTTP_400_BAD_REQUEST, f"limit must be in 1..{max_limit}." ) - call = AgentCall.get(call_id) + call = await AgentCall.get(call_id) if call.live_status == AgentCallStatus.RUNNING: response.headers["Cache-Control"] = "no-store" else: @@ -367,7 +367,7 @@ async def stream_transcript( @router.get("/files", response_model=AgentCallFiles, response_model_by_alias=True) async def list_files(call_id: str) -> AgentCallFiles: - return reads.get_agent_call_files(call_id) + return await reads.get_agent_call_files(call_id) @router.get("/files/{file_name:path}") async def get_file( @@ -375,7 +375,7 @@ async def get_file( file_name: str, disposition: Literal["inline", "attachment"] = "inline", ) -> FileResponse: - call = AgentCall.get(call_id) + call = await AgentCall.get(call_id) resolved = call.get_file_path(file_name) if not resolved: raise HTTPException(status.HTTP_404_NOT_FOUND, "File not found for this call.") @@ -411,22 +411,22 @@ def _get_subject_routes( subject_type = subject_class.subject_type router = APIRouter(prefix=f"/{subject_type}", tags=[f"{cls.name}:{subject_type}"]) - def board(account_id: str | None) -> SubjectList: + async def board(account_id: str | None) -> SubjectList: return SubjectList( rows=[ SubjectRow( summary=summary, - status=reads.get_subject_status(subject_type, summary.id), + status=await reads.get_subject_status(subject_type, summary.id), ) - for summary in subject_class.list_summaries(account_id) + for summary in await subject_class.list_summaries(account_id) ] ) async def subject_response(subject_id: str) -> SubjectResponse | None: - subject = subject_class.get_for_subject_id(subject_id) + subject = await subject_class.get_for_subject_id(subject_id) if subject is None: return - return reads.get_subject_response( + return await reads.get_subject_response( subject_type, subject_id, summary=subject.get_summary(), @@ -435,7 +435,7 @@ async def subject_response(subject_id: str) -> SubjectResponse | None: @router.get("", response_model=SubjectList, response_model_by_alias=True) async def list_subjects() -> SubjectList: - return board(current_account_id.get()) + return await board(current_account_id.get()) # ``/stream`` before ``/{subject_id}`` so the literal path wins over the id matcher. @router.get("/stream", response_class=StreamingResponse) @@ -444,8 +444,8 @@ async def stream_board(engine: EngineDep) -> StreamingResponse: account_id = current_account_id.get() async def snapshot() -> SubjectList: - with session_scope(engine): - return board(account_id) + async with session_scope(engine): + return await board(account_id) return StreamingResponse( stream(snapshot), media_type="text/event-stream", headers=SSE_HEADERS @@ -457,7 +457,7 @@ async def snapshot() -> SubjectList: @router.get("/{subject_id:path}/stream", response_class=StreamingResponse) async def stream_subject(subject_id: str, engine: EngineDep) -> StreamingResponse: async def snapshot() -> SubjectResponse | None: - with session_scope(engine): + async with session_scope(engine): return await subject_response(subject_id) return StreamingResponse( @@ -483,7 +483,7 @@ async def on_startup(cls) -> None: boot.""" @classmethod - def record_event( + async def record_event( cls, *, type: str, @@ -494,7 +494,7 @@ def record_event( app automatically. Apps record through here so the ``Event`` model stays a platform internal. ``type`` is the milestone's own word ("merged") — the feed reads it as one, so an app writes no rendering.""" - Event.emit( + await Event.emit( type=type, subject=subject.identity if subject else None, label=subject.label if subject else None, diff --git a/backend/druks/apps/fetcher.py b/backend/druks/apps/fetcher.py index 2c8a6946..a3ff970a 100644 --- a/backend/druks/apps/fetcher.py +++ b/backend/druks/apps/fetcher.py @@ -17,7 +17,7 @@ async def fetch_file(*, repo: str, path: str) -> str | None: if time.time() - cache.stat().st_mtime < _TTL_SECONDS: return cache.read_text() or None - github = get_github_client() + github = await get_github_client() try: body = await github.get_file_content(repo, path) diff --git a/backend/druks/bootstrap.py b/backend/druks/bootstrap.py index 227426f2..ee4ce33e 100644 --- a/backend/druks/bootstrap.py +++ b/backend/druks/bootstrap.py @@ -1,8 +1,8 @@ from sqlalchemy import select +from sqlalchemy.orm import Session from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.accounts.models import Account -from druks.database import get_session from druks.harnesses.registry import get_harnesses from druks.user_settings.models import HarnessSettings @@ -19,7 +19,7 @@ def seed_harnesses(engine) -> None: # Give every registered harness a config row at its shipped defaults. A # harness added later is seeded on the deploy that adds it; existing rows # keep whatever the operator tuned. - with get_session(engine) as session: + with Session(engine) as session: existing = set(session.execute(select(HarnessSettings.name)).scalars()) for harness in get_harnesses(): if harness.name not in existing: @@ -36,7 +36,7 @@ def seed_harnesses(engine) -> None: def seed_system_account(engine) -> None: # Owns every run nobody asked for: crons, background work. - with get_session(engine) as session: + with Session(engine) as session: if not session.get(Account, SYSTEM_ACCOUNT_ID): session.add(Account(id=SYSTEM_ACCOUNT_ID, username="system")) session.commit() diff --git a/backend/druks/browser/login.py b/backend/druks/browser/login.py index 32caf91e..59c7cf01 100644 --- a/backend/druks/browser/login.py +++ b/backend/druks/browser/login.py @@ -105,13 +105,13 @@ async def save(self) -> StoredBrowserSession: """Store what the operator logged into as the session's payload, then tear the window down. A login always captures a profile, so a session imported as storage_state becomes a profile here.""" - session = StoredBrowserSession.get_for_name(self.session_name) + session = await StoredBrowserSession.get_for_name(self.session_name) if session: try: async with sandbox_client.attach(host_id=self.host_id) as browser: payload = await _export(browser, session.name) session.payload_format = BrowserSessionPayloadFormat.PROFILE_DIR.value - session.store_payload(payload) + await session.store_payload(payload) return session finally: await self._close() diff --git a/backend/druks/browser/models.py b/backend/druks/browser/models.py index b5c99c37..beb211f2 100644 --- a/backend/druks/browser/models.py +++ b/backend/druks/browser/models.py @@ -38,7 +38,7 @@ class StoredBrowserSession(Base, Uuid7Pk): last_used_at: Mapped[datetime | None] = mapped_column(default=None) @classmethod - def get_or_create( + async def get_or_create( cls, *, name: str, @@ -49,42 +49,43 @@ def get_or_create( """Concurrency-safe lookup-or-create: two first actions racing on the same session both INSERT with ON CONFLICT DO NOTHING, then converge on the one row through the name lookup.""" - browser_session = cls.get_for_name(name) + browser_session = await cls.get_for_name(name) if browser_session: return browser_session session = db_session() - session.execute( + await session.execute( insert(cls) .values(name=name, payload_format=payload_format.value, site=site, status=status.value) .on_conflict_do_nothing(index_elements=["name"]) ) - return session.scalars(select(cls).where(cls.name == name)).one() + return (await session.scalars(select(cls).where(cls.name == name))).one() @classmethod - def list_all(cls): - return list(db_session().scalars(select(cls).order_by(cls.name))) + async def list_all(cls): + return list(await db_session().scalars(select(cls).order_by(cls.name))) @classmethod - def get_for_name(cls, name: str): - return db_session().scalar(select(cls).where(cls.name == name)) + async def get_for_name(cls, name: str): + return await db_session().scalar(select(cls).where(cls.name == name)) - def mark_stale(self) -> None: + async def mark_stale(self) -> None: self.status = BrowserSessionStatus.STALE.value - db_session().flush() + await db_session().flush() - def mark_used(self) -> None: + async def mark_used(self) -> None: self.last_used_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def store_payload(self, payload: bytes) -> None: + async def store_payload(self, payload: bytes) -> None: self.payload = payload # type: ignore[assignment] — the column takes plaintext in, hands SecretBytes back self.status = BrowserSessionStatus.READY.value self.last_refreshed_at = Base.utc_now() - db_session().flush() + await db_session().flush() # Assignment holds the plaintext; a read must always hand back the - # encrypted column's SecretBytes, so the next access reloads. - db_session().expire(self, ["payload"]) + # encrypted column's SecretBytes, so reload the column now — an expired + # attribute can't lazy-load under the async session. + await db_session().refresh(self, ["payload"]) - def delete(self) -> None: - db_session().delete(self) - db_session().flush() + async def delete(self) -> None: + await db_session().delete(self) + await db_session().flush() diff --git a/backend/druks/browser/routes.py b/backend/druks/browser/routes.py index 0028f903..93720d0c 100644 --- a/backend/druks/browser/routes.py +++ b/backend/druks/browser/routes.py @@ -25,7 +25,7 @@ @router.get("", response_model=list[BrowserSessionResponse]) async def list_browser_sessions(account: Account = Depends(current_account)): - rows = {row.name: row for row in StoredBrowserSession.list_all()} + rows = {row.name: row for row in await StoredBrowserSession.list_all()} sessions = [] for declaration in browser_sessions.all(): try: @@ -66,7 +66,7 @@ async def upload_state( if declaration := browser_sessions.get(name): if declaration.anonymous: raise exceptions.BrowserSessionAnonymousError(name) - row = declaration.get_or_create_row() + row = await declaration.get_or_create_row() payload = bytearray() async for chunk in request.stream(): payload.extend(chunk) @@ -78,7 +78,7 @@ async def upload_state( if len(payload) >= PAYLOAD_WARNING_BYTES: logger.warning("Browser session %s received a %d-byte payload.", name, len(payload)) row.payload_format = payload_format.value - row.store_payload(bytes(payload)) + await row.store_payload(bytes(payload)) return raise exceptions.BrowserSessionUnknownError(name) @@ -91,7 +91,7 @@ async def open_login_window( if declaration := browser_sessions.get(name): if declaration.anonymous: raise exceptions.BrowserSessionAnonymousError(name) - await LoginWindow.open(declaration.get_or_create_row()) + await LoginWindow.open(await declaration.get_or_create_row()) return raise exceptions.BrowserSessionUnknownError(name) @@ -102,7 +102,7 @@ async def login_window_socket(websocket: WebSocket, name: str) -> None: await websocket.close(code=1008) return try: - with session_scope(websocket.app.state.engine): + async with session_scope(websocket.app.state.engine): await require_operator(websocket) window = await LoginWindow.get_for_session(name) except (HTTPException, exceptions.BrowserApiError): @@ -137,7 +137,7 @@ async def delete_browser_session( name: str, account: Account = Depends(current_session_account), ) -> None: - if row := StoredBrowserSession.get_for_name(name): - row.delete() + if row := await StoredBrowserSession.get_for_name(name): + await row.delete() return raise exceptions.BrowserSessionUnknownError(name) diff --git a/backend/druks/browser/sessions.py b/backend/druks/browser/sessions.py index 8333c4dc..4dd5f075 100644 --- a/backend/druks/browser/sessions.py +++ b/backend/druks/browser/sessions.py @@ -77,7 +77,7 @@ async def cdp(self): block. The browser lives in its own container on the druks box and dies with the block; a persisting session is exported and stored back first.""" - row = self.get_or_create_row() if self.anonymous else self._ready_row() + row = await (self.get_or_create_row() if self.anonymous else self._ready_row()) writer_token = await acquire_writer_lock(row.id) if self.persist else "" try: settings = load_settings() @@ -87,7 +87,7 @@ async def cdp(self): ) as browser: await seed_state(browser, row) await self._launch(browser) - row.mark_used() + await row.mark_used() listener = await browser.forward_local_port(CDP_PORT) try: yield f"http://127.0.0.1:{listener.get_port()}" @@ -100,7 +100,7 @@ async def cdp(self): listener.close() if self.persist: row.payload_format = BrowserSessionPayloadFormat.PROFILE_DIR.value - row.store_payload(await self._export(browser)) + await row.store_payload(await self._export(browser)) finally: if writer_token: await release_writer_lock(row.id, writer_token) @@ -123,20 +123,20 @@ async def playwright(self): finally: await connection.close() - def get_or_create_row(self) -> StoredBrowserSession: + async def get_or_create_row(self) -> StoredBrowserSession: """The declaration's stored half, written by the first action that needs it — a borrow, a login-window open, or a state import. Until then the declaration alone puts the session in the pane, wanting a login.""" - return StoredBrowserSession.get_or_create( + return await StoredBrowserSession.get_or_create( name=self.name, payload_format=BrowserSessionPayloadFormat.PROFILE_DIR, site=self.site, status=self.initial_status, ) - def _ready_row(self) -> StoredBrowserSession: - row = self.get_or_create_row() + async def _ready_row(self) -> StoredBrowserSession: + row = await self.get_or_create_row() if row.status != BrowserSessionStatus.READY.value: raise BrowserSessionNotReadyError(self.name, row.status) return row diff --git a/backend/druks/browser/subscribers.py b/backend/druks/browser/subscribers.py index ad96559c..0f2148bd 100644 --- a/backend/druks/browser/subscribers.py +++ b/backend/druks/browser/subscribers.py @@ -10,6 +10,6 @@ async def signed_out_session_goes_stale(*, session_name: str, **_: object) -> No # session goes stale — the pane shows it and refuses borrows until a re-login. # An anonymous session has no login to go stale: the run still fails, the # row stays anonymous. - row = StoredBrowserSession.get_for_name(session_name) + row = await StoredBrowserSession.get_for_name(session_name) if row.status != BrowserSessionStatus.ANONYMOUS.value: - row.mark_stale() + await row.mark_stale() diff --git a/backend/druks/contrib/review/app.py b/backend/druks/contrib/review/app.py index a4dd5660..aec6598a 100644 --- a/backend/druks/contrib/review/app.py +++ b/backend/druks/contrib/review/app.py @@ -6,11 +6,11 @@ from druks.doctor import CheckResult -def check_review_identity() -> CheckResult: +async def check_review_identity() -> CheckResult: """Set or unset, both healthy: an empty pair is comment mode by design. A half-configured pair is the settings clean's failure (``review:settings``), not this check's.""" - settings = Review.settings() + settings = await Review.settings() if settings.app_id and settings.private_key: return CheckResult( name="identity", ok=True, detail="set — reviews approve as the distinct App" diff --git a/backend/druks/contrib/review/datastructures.py b/backend/druks/contrib/review/datastructures.py index 81fae076..c99812ed 100644 --- a/backend/druks/contrib/review/datastructures.py +++ b/backend/druks/contrib/review/datastructures.py @@ -14,7 +14,7 @@ def get(cls, repo: str, number: int) -> Self: return cls(id=f"{repo}#{number}") @classmethod - def get_for_subject_id(cls, subject_id: str) -> Self | None: + async def get_for_subject_id(cls, subject_id: str) -> Self | None: # Ids reach the read side as free text off a URL, so a shape that names no # pull request is a miss rather than a crashed read. repo, _, number = subject_id.partition("#") @@ -45,7 +45,7 @@ def get_summary(self) -> ReviewSummary: ) @classmethod - def list_summaries(cls, account_id: str | None) -> list[ReviewSummary]: + async def list_summaries(cls, account_id: str | None) -> list[ReviewSummary]: """The reviews still going or stopped on a failure. A finished one lives on its pull request, so it leaves the board as soon as it has something to show there.""" - return [pull_request.get_summary() for pull_request in cls.list_open()] + return [pull_request.get_summary() for pull_request in await cls.list_open()] diff --git a/backend/druks/contrib/review/github.py b/backend/druks/contrib/review/github.py index ca30e713..3e444951 100644 --- a/backend/druks/contrib/review/github.py +++ b/backend/druks/contrib/review/github.py @@ -18,8 +18,8 @@ class ReviewActor: mode: Literal["approve", "comment"] -def get_review_actor() -> ReviewActor: - settings = Review.settings() +async def get_review_actor() -> ReviewActor: + settings = await Review.settings() if settings.app_id and settings.private_key: # Only a complete pair selects the distinct identity — a half-configured # one (flagged by clean()) still borrows the operator client below. @@ -31,4 +31,4 @@ def get_review_actor() -> ReviewActor: ), mode="approve", ) - return ReviewActor(client=get_github_client(), mode="comment") + return ReviewActor(client=await get_github_client(), mode="comment") diff --git a/backend/druks/contrib/review/routes.py b/backend/druks/contrib/review/routes.py index 7cccc3f4..6216e76d 100644 --- a/backend/druks/contrib/review/routes.py +++ b/backend/druks/contrib/review/routes.py @@ -46,7 +46,7 @@ async def request_review( account: Account = Depends(current_account), ) -> str: """Start a pull request review.""" - if not ProjectRepo.get_for_repo(repo): + if not await ProjectRepo.get_for_repo(repo): raise HTTPException( status.HTTP_404_NOT_FOUND, f"{repo} is not a registered project repo — add it to a project first", diff --git a/backend/druks/contrib/review/subscribers.py b/backend/druks/contrib/review/subscribers.py index 658dfbde..bd3867b1 100644 --- a/backend/druks/contrib/review/subscribers.py +++ b/backend/druks/contrib/review/subscribers.py @@ -8,9 +8,9 @@ async def mention_asks_for_a_review(*, repo: str, pr_number: int, payload: dict) -> None: """Addressing the review actor asks it to review that pull request, and only someone who writes to the repo may ask — a review is the account's to spend.""" - handle = await get_review_actor().client.get_mention_handle() + handle = await (await get_review_actor()).client.get_mention_handle() is_mentioned = handle and f"@{handle}".casefold() in payload["body"].casefold() - if is_mentioned and ProjectRepo.get_for_repo(repo): + if is_mentioned and await ProjectRepo.get_for_repo(repo): await PullRequestReview.dispatch( repo=repo, pr_number=pr_number, requested_by=payload["author"] ) diff --git a/backend/druks/contrib/review/workflows.py b/backend/druks/contrib/review/workflows.py index 7c9fb477..7eff7434 100644 --- a/backend/druks/contrib/review/workflows.py +++ b/backend/druks/contrib/review/workflows.py @@ -43,10 +43,10 @@ async def dispatch(cls, *, repo: str, pr_number: int, requested_by: str) -> str: # Even a distinct review identity clones alongside the operator App, so # resolve the operator identity before the start spends a run and # provisions a VM — the raising lookup surfaces the actionable error. - ServiceIdentity.get(GITHUB) + await ServiceIdentity.get(GITHUB) # Attribution follows the requester when druks knows them by that name; a # review asked for by someone with no account runs as the system's. - account = Account.get_for_username(requested_by) + account = await Account.get_for_username(requested_by) return await cls.start( subject=PullRequest.get(repo, pr_number), account_id=account.id if account else None, @@ -61,8 +61,8 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: # The token is the review actor's: it authenticates the clone and ``gh``, so the # review is authored under that identity. Siblings stay uncloned — the reviewer # clones the ones it opens, into a directory that must exist for the grant to hold. - repo = self.subject.repo - github_token = await get_review_actor().client.token_for_repo(repo) + repo = (await self.subject).repo + github_token = await (await get_review_actor()).client.token_for_repo(repo) await sandbox.write_secret( secret=github_token, remote=get_github_token_remote_path(sandbox.ssh_username) ) @@ -80,9 +80,9 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: } async def get_prompt_context(self, **context: Any) -> dict[str, Any]: - target = ProjectRepo.get_for_repo(self.subject.repo, raise_on_missing=True) + target = await ProjectRepo.get_for_repo((await self.subject).repo, raise_on_missing=True) return { - "siblings": target.siblings(), - "review_mode": get_review_actor().mode, + "siblings": await target.siblings(), + "review_mode": (await get_review_actor()).mode, **await super().get_prompt_context(**context), } diff --git a/backend/druks/contrib/ship/app.py b/backend/druks/contrib/ship/app.py index cc8d7e21..c6cd79f4 100644 --- a/backend/druks/contrib/ship/app.py +++ b/backend/druks/contrib/ship/app.py @@ -29,14 +29,14 @@ } -def check_tracker_identity() -> CheckResult: +async def check_tracker_identity() -> CheckResult: """Whether the selected tracker's identity is connected. Trackerless is a choice, not a fault; a selected-but-unconnected tracker is pending setup.""" - settings = Ship.settings() + settings = await Ship.settings() if settings.tracker == "none": return CheckResult(name="tracker", ok=True, detail="trackerless by choice") service = {"linear": services.Linear, "jira": services.Jira}[settings.tracker] - if service.is_connected(): + if await service.is_connected(): return CheckResult(name="tracker", ok=True, detail=f"{settings.tracker} connected") return CheckResult( name="tracker", @@ -110,24 +110,24 @@ def trigger_status(self) -> str: checks = [check_tracker_identity] @classmethod - def get_tracker(cls, source: str | None = None) -> Tracker | None: + async def get_tracker(cls, source: str | None = None) -> Tracker | None: """The selected tracker, once its service identity is connected; None when the installation runs trackerless or the identity is missing. Pass a ``source`` to get it only when that source is the selected one — a work item syncs only to the tracker that owns it.""" - settings = cls.settings() + settings = await cls.settings() if source is not None and source != settings.tracker: return try: if settings.tracker == "linear": - row = services.Linear.get() + row = await services.Linear.get() return Linear( api_key=row.secrets["api_key"], backlog_status=settings.linear_resting_status, trigger_status=settings.trigger_status, ) if settings.tracker == "jira": - row = services.Jira.get() + row = await services.Jira.get() return Jira( base_url=row.identity["base_url"], email=row.identity["email"], diff --git a/backend/druks/contrib/ship/models.py b/backend/druks/contrib/ship/models.py index 0491fc14..9a906237 100644 --- a/backend/druks/contrib/ship/models.py +++ b/backend/druks/contrib/ship/models.py @@ -35,19 +35,21 @@ class Project(Base): ) @classmethod - def create(cls, *, name: str) -> "Project": + async def create(cls, *, name: str) -> "Project": session = db_session() - project = cls(name=name) + # Seed the collection as loaded-empty: a fresh project has no repos, and + # the summary read right after flush must not trigger a lazy load. + project = cls(name=name, repos=[]) session.add(project) - session.flush() + await session.flush() return project @classmethod - def get(cls, project_id: int) -> "Project | None": - return db_session().get(cls, project_id) + async def get(cls, project_id: int) -> "Project | None": + return await db_session().get(cls, project_id) @classmethod - def get_for_repo(cls, full_name: str) -> "Project | None": + async def get_for_repo(cls, full_name: str) -> "Project | None": """Lookup the Project that owns ``full_name`` (e.g. ``clawhaven/acme-app``). Returns None when the repo isn't bound to any project yet — the @@ -59,7 +61,7 @@ def get_for_repo(cls, full_name: str) -> "Project | None": .where(func.lower(ProjectRepo.full_name) == full_name.lower()) .limit(1) ) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() class ProjectRepo(StoredSubject): @@ -79,7 +81,7 @@ class ProjectRepo(StoredSubject): project: Mapped[Project] = relationship(back_populates="repos", lazy="joined") @classmethod - def create( + async def create( cls, *, project_id: int, @@ -89,19 +91,19 @@ def create( session = db_session() row = cls(project_id=project_id, full_name=full_name, purpose=purpose) session.add(row) - session.flush() + await session.flush() return row @classmethod - def get(cls, repo_id: int) -> "ProjectRepo | None": - return db_session().get(cls, repo_id) + async def get(cls, repo_id: int) -> "ProjectRepo | None": + return await db_session().get(cls, repo_id) @classmethod - def get_in_project(cls, *, project_id: int, repo_id: int) -> "ProjectRepo | None": + async def get_in_project(cls, *, project_id: int, repo_id: int) -> "ProjectRepo | None": # Scoped lookup for the nested /projects/{project_id}/repos/{repo_id} routes: # a repo reached through the wrong project's URL is a miss, not a hit to reject. stmt = select(cls).where(cls.id == repo_id, cls.project_id == project_id).limit(1) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() def get_label(self) -> str: return self.full_name @@ -110,25 +112,31 @@ def get_summary(self) -> "ProjectRepoSummary": return ProjectRepoSummary.model_validate(self) @classmethod - def list_summaries(cls, account_id: str | None) -> list["ProjectRepoSummary"]: + async def list_summaries(cls, account_id: str | None) -> list["ProjectRepoSummary"]: # A repo is registered, not transient, so the board is all of them by name. stmt = select(cls).order_by(cls.full_name) - return [repo.get_summary() for repo in db_session().scalars(stmt)] - - def siblings(self) -> list["ProjectRepo"]: - return [repo for repo in self.project.repos if repo.full_name != self.full_name] + return [repo.get_summary() for repo in await db_session().scalars(stmt)] + + async def siblings(self) -> list["ProjectRepo"]: + # A fresh query, not ``self.project.repos``: the loaded collection goes + # stale when a repo is registered through its FK in the same session. + stmt = select(ProjectRepo).where( + ProjectRepo.project_id == self.project_id, + ProjectRepo.full_name != self.full_name, + ) + return list(await db_session().scalars(stmt)) @property def effective_profile(self) -> dict[str, Any]: # {} until the repo profiler has run — an unprofiled repo is a normal state. return self.profile.get("effective") or {} - def set_profile(self, *, baseline: dict[str, Any], effective: dict[str, Any]) -> None: + async def set_profile(self, *, baseline: dict[str, Any], effective: dict[str, Any]) -> None: self.profile = {"baseline": baseline, "effective": effective} - db_session().flush() + await db_session().flush() @classmethod - def get_for_name(cls, name: str) -> "ProjectRepo | None": + async def get_for_name(cls, name: str) -> "ProjectRepo | None": """Match a ticket signal against the bare repo name. Convention: a tracker project name (Linear) or a label names the @@ -141,14 +149,14 @@ def get_for_name(cls, name: str) -> "ProjectRepo | None": return # SQLite-friendly bare-name suffix match. stmt = select(cls).where(func.lower(cls.full_name).like(f"%/{target}")).limit(1) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() @classmethod - def get_for_repo( + async def get_for_repo( cls, full_name: str, *, raise_on_missing: bool = False ) -> "ProjectRepo | None": stmt = select(cls).where(func.lower(cls.full_name) == full_name.lower()).limit(1) - repo = db_session().scalars(stmt).first() + repo = (await db_session().scalars(stmt)).first() if raise_on_missing and not repo: # A run's stored repo name can outlive its registration — a rename or # a GitHub transfer leaves the old full_name behind, and this matches @@ -160,7 +168,7 @@ def get_for_repo( return repo @classmethod - def lookup( + async def lookup( cls, *, project_name: str | None, @@ -175,7 +183,7 @@ def lookup( """ for name in (project_name, *labels): if name: - row = cls.get_for_name(name) + row = await cls.get_for_name(name) if row: return row return @@ -228,16 +236,16 @@ def get_summary(self) -> WorkItemSummary: return WorkItemSummary.model_validate(self) @classmethod - def list_summaries(cls, account_id: str | None) -> list[WorkItemSummary]: + async def list_summaries(cls, account_id: str | None) -> list[WorkItemSummary]: # Where a run stands colours the row; it never decides whether the row is # here. The 500 most-recent cover it; paginate if a board outgrows it. stmt = ( select(cls).where(cls.resolution.is_(None)).order_by(cls.updated_at.desc()).limit(500) ) - return [item.get_summary() for item in db_session().scalars(stmt)] + return [item.get_summary() for item in await db_session().scalars(stmt)] @classmethod - def create( + async def create( cls, *, project_id: int, @@ -257,15 +265,15 @@ def create( repo=repo, ) session.add(item) - session.flush() + await session.flush() return item @classmethod - def get(cls, work_item_id: int) -> "WorkItem | None": - return db_session().get(cls, work_item_id) + async def get(cls, work_item_id: int) -> "WorkItem | None": + return await db_session().get(cls, work_item_id) @classmethod - def get_for_pr( + async def get_for_pr( cls, *, repo: str, pr_number: int | None, branch: str | None = None ) -> "WorkItem | None": """The item a pull request belongs to. Its number identifies it; the head @@ -277,45 +285,45 @@ def get_for_pr( .order_by(cls.updated_at.desc()) .limit(1) ) - found = db_session().scalars(stmt).first() + found = (await db_session().scalars(stmt)).first() if found: return found - return cls.get_for_branch(repo=repo, branch=branch) if branch else None + return await cls.get_for_branch(repo=repo, branch=branch) if branch else None @classmethod - def get_for_branch(cls, *, repo: str, branch: str) -> "WorkItem | None": + async def get_for_branch(cls, *, repo: str, branch: str) -> "WorkItem | None": stmt = ( select(cls) .where(func.lower(cls.repo) == repo.lower(), cls.branch == branch) .order_by(cls.updated_at.desc()) .limit(1) ) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() - def start_attempt(self) -> None: + async def start_attempt(self) -> None: self.branch = None self.pr_number = None self.resolution = None self.resolved_at = None self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def resolve(self, *, merged: bool, at: datetime) -> None: + async def resolve(self, *, merged: bool, at: datetime) -> None: # cycle: the app imports this module at file scope. import druks.contrib.ship.app as ship_app self.resolution = "merged" if merged else "closed" self.resolved_at = at self.updated_at = Base.utc_now() - ship_app.Ship.record_event(type=self.resolution, subject=self) - db_session().flush() + await ship_app.Ship.record_event(type=self.resolution, subject=self) + await db_session().flush() async def ship(self) -> None: # A build parked on the operator's review is stranded by their merge; a running # one converges on its own, its merge step finding the PR already closed. from druks.contrib.ship.workflows import Build - build = self.get_status(workflow=Build) + build = await self.get_status(workflow=Build) if build.is_parked: await Build.cancel(self, failure="pr merged while parked") await self.set_ticket_status(TicketStatus.DONE) @@ -327,16 +335,16 @@ async def close_external(self) -> None: from druks.contrib.ship.workflows import Build await Build.cancel(self, failure="pr closed without merge") - db_session().flush() + await db_session().flush() try: if (await RepoPolicy.resolve(self.repo)).delete_branch: - await get_github_client().delete_branch(self.repo, self.branch) + await (await get_github_client()).delete_branch(self.repo, self.branch) except Exception: # noqa: BLE001 — cleanup only logger.warning("Skipped branch cleanup for %s.", self.repo, exc_info=True) await self.set_ticket_status(TicketStatus.BACKLOG) @classmethod - def get_for_ticket_key( + async def get_for_ticket_key( cls, *, source: str, @@ -345,28 +353,28 @@ def get_for_ticket_key( """The item a ticket names in its tracker — (source, ticket_key) is the row's identity.""" stmt = select(cls).where(cls.source == source, cls.ticket_key == ticket_key).limit(1) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() @classmethod - def list_recent(cls, *, limit: int = 50, offset: int = 0) -> list["WorkItem"]: + async def list_recent(cls, *, limit: int = 50, offset: int = 0) -> list["WorkItem"]: stmt = select(cls).order_by(cls.updated_at.desc()).limit(limit).offset(offset) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def list_handoff(cls, *, limit: int = 10) -> list["WorkItem"]: + async def list_handoff(cls, *, limit: int = 10) -> list["WorkItem"]: stmt = ( select(cls) .where(cls.resolved_at.is_not(None)) .order_by(cls.resolved_at.desc()) .limit(limit) ) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) async def set_ticket_status(self, status: TicketStatus) -> None: # Lazy: the Ship app imports this module, so it can't be imported at top. import druks.contrib.ship.app as ship_app - tracker = ship_app.Ship.get_tracker(self.source) + tracker = await ship_app.Ship.get_tracker(self.source) # No tracker means nothing to sync: a github item, a source the operator # has switched away from, or credentials not set yet. if not tracker: @@ -384,7 +392,7 @@ async def set_ticket_status(self, status: TicketStatus) -> None: exc_info=True, ) - def update( + async def update( self, *, title: str = _KEEP, @@ -404,4 +412,4 @@ def update( if project_id is not _KEEP: self.project_id = project_id self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() diff --git a/backend/druks/contrib/ship/routes.py b/backend/druks/contrib/ship/routes.py index f80172f1..84ce3835 100644 --- a/backend/druks/contrib/ship/routes.py +++ b/backend/druks/contrib/ship/routes.py @@ -37,7 +37,7 @@ @projects_router.get("", response_model=ProjectsResponse, response_model_by_alias=True) async def list_projects() -> ProjectsResponse: - rows = list(db_session().scalars(select(Project).order_by(Project.name))) + rows = list(await db_session().scalars(select(Project).order_by(Project.name))) return ProjectsResponse(projects=[ProjectSummary.model_validate(p) for p in rows]) @@ -51,7 +51,7 @@ async def create_project(body: CreateProjectRequest) -> ProjectSummary: name = body.name.strip() if not name: raise HTTPException(status.HTTP_400_BAD_REQUEST, "name is required") - project = Project.create(name=name) + project = await Project.create(name=name) return ProjectSummary.model_validate(project) @@ -72,7 +72,7 @@ async def list_github_repos( ), ), ) -> GitHubReposResponse: - github = get_github_client() + github = await get_github_client() resolved = (owner or "").strip() if resolved: owners: tuple[str, ...] = (resolved,) @@ -102,7 +102,7 @@ async def list_github_repos( response_model_by_alias=True, ) async def get_project(project_id: int) -> ProjectSummary: - project = Project.get(project_id) + project = await Project.get(project_id) if not project: raise HTTPException(status.HTTP_404_NOT_FOUND, "project not found") return ProjectSummary.model_validate(project) @@ -117,7 +117,7 @@ async def update_project( project_id: int, name: str | None = Body(default=None, embed=True), ) -> ProjectSummary: - project = Project.get(project_id) + project = await Project.get(project_id) if not project: raise HTTPException(status.HTTP_404_NOT_FOUND, "project not found") if name is not None: @@ -125,7 +125,7 @@ async def update_project( if not name: raise HTTPException(status.HTTP_400_BAD_REQUEST, "name cannot be empty") project.name = name - db_session().flush() + await db_session().flush() return ProjectSummary.model_validate(project) @@ -136,12 +136,12 @@ async def delete_project(project_id: int) -> None: cascade is an explicit child delete in the same session before the project (and its repo ``delete-orphan`` cascade) is deleted.""" session = db_session() - project = Project.get(project_id) + project = await Project.get(project_id) if not project: raise HTTPException(status.HTTP_404_NOT_FOUND, "project not found") - session.execute(delete(WorkItem).where(WorkItem.project_id == project_id)) - session.delete(project) - session.flush() + await session.execute(delete(WorkItem).where(WorkItem.project_id == project_id)) + await session.delete(project) + await session.flush() @projects_router.post( @@ -154,7 +154,7 @@ async def add_project_repo( project_id: int, body: AddProjectRepoRequest, ) -> ProjectRepoSummary: - project = Project.get(project_id) + project = await Project.get(project_id) if not project: raise HTTPException(status.HTTP_404_NOT_FOUND, "project not found") full_name = body.full_name.strip() @@ -163,13 +163,13 @@ async def add_project_repo( status.HTTP_400_BAD_REQUEST, "fullName must be 'owner/name'", ) - repo = ProjectRepo.create( + repo = await ProjectRepo.create( project_id=project.id, full_name=full_name, purpose=body.purpose, ) # Re-bind matching work items that still point at a different project. - db_session().execute( + await db_session().execute( update(WorkItem) .where( func.lower(WorkItem.repo) == full_name.lower(), @@ -196,12 +196,12 @@ async def update_project_repo( repo_id: int, purpose: str | None = Body(default=None, embed=True), ) -> ProjectRepoSummary: - row = ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) + row = await ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) if not row: raise HTTPException(status.HTTP_404_NOT_FOUND, "repo not found") if purpose is not None: row.purpose = purpose.strip() or None - db_session().flush() + await db_session().flush() return ProjectRepoSummary.model_validate(row) @@ -211,7 +211,7 @@ async def update_project_repo( response_model_by_alias=True, ) async def profile_project_repo(project_id: int, repo_id: int) -> ProjectRepoSummary: - row = ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) + row = await ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) if not row: raise HTTPException(status.HTTP_404_NOT_FOUND, "repo not found") # Profile is subject-unique: dispatch() returns the live run when one is already @@ -225,12 +225,12 @@ async def profile_project_repo(project_id: int, repo_id: int) -> ProjectRepoSumm status_code=status.HTTP_204_NO_CONTENT, ) async def delete_project_repo(project_id: int, repo_id: int) -> None: - row = ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) + row = await ProjectRepo.get_in_project(project_id=project_id, repo_id=repo_id) if not row: raise HTTPException(status.HTTP_404_NOT_FOUND, "repo not found") session = db_session() - session.delete(row) - session.flush() + await session.delete(row) + await session.flush() # /api/ship/work-items WorkItem CRUD @@ -256,7 +256,7 @@ async def list_work_items_history( response.headers["Cache-Control"] = "no-store" clamped = max(1, min(limit, _HISTORY_MAX_LIMIT)) # Recent history: the PRs GitHub has resolved, its verdict newest first. - items = [DashboardItem.model_validate(wi) for wi in WorkItem.list_handoff(limit=clamped)] + items = [DashboardItem.model_validate(wi) for wi in await WorkItem.list_handoff(limit=clamped)] return WorkItemsHistoryResponse(items=items) @@ -277,7 +277,7 @@ async def start_work_item( """Move the tracker ticket to the configured trigger status; webhook intake then opens the build. No run exists yet when this returns — poll list_open_subjects while waiting for it.""" - tracker = Ship.get_tracker() + tracker = await Ship.get_tracker() if not tracker: raise TrackerNotConfigured() async with tracker: diff --git a/backend/druks/contrib/ship/subscribers.py b/backend/druks/contrib/ship/subscribers.py index 9a79ea3c..ab3014a8 100644 --- a/backend/druks/contrib/ship/subscribers.py +++ b/backend/druks/contrib/ship/subscribers.py @@ -10,14 +10,14 @@ @subscribe(WorkflowEvent.SCHEDULED, workflow=Build) async def new_build_claims_the_item(*, subject: WorkItem, **_: object) -> None: - subject.start_attempt() + await subject.start_attempt() @subscribe(WorkflowEvent.CANCELLED, workflow=Build) async def cancelled_build_settles_the_item(*, subject: WorkItem, **_: object) -> None: """An operator cancellation explicitly abandons the work item.""" if not subject.resolution: - subject.resolve(merged=False, at=Base.utc_now()) + await subject.resolve(merged=False, at=Base.utc_now()) @subscribe("pr.opened", workflow=Build) @@ -26,7 +26,7 @@ async def pr_open_mirrors_onto_item( ) -> None: # The implementer's provisioned PR + branch, mirrored onto the work item — # the read side (board links, webhook routing by repo+PR) keys off them. - subject.update(pr_number=pr_number, branch=branch) + await subject.update(pr_number=pr_number, branch=branch) @subscribe(WorkflowEvent.RUNNING, workflow=Build) @@ -46,7 +46,7 @@ async def policy_push_reprofiles_the_repo(*, repo: str, paths: list, **_: object # The operator edited the repo's build policy — re-apply it over the # profiled baseline. if ".druks/ship/config.yml" in paths: - project_repo = ProjectRepo.get_for_repo(repo) + project_repo = await ProjectRepo.get_for_repo(repo) if project_repo: await Profile.dispatch(project_repo, refresh_only=True) @@ -54,10 +54,10 @@ async def policy_push_reprofiles_the_repo(*, repo: str, paths: list, **_: object @subscribe("pr.review_submitted") async def pr_review_answers_the_gate(*, repo: str, pr_number: int, payload: dict) -> None: - item = WorkItem.get_for_pr(repo=repo, pr_number=pr_number, branch=payload["branch"]) + item = await WorkItem.get_for_pr(repo=repo, pr_number=pr_number, branch=payload["branch"]) if not item: return - status = item.get_status(workflow=Build) + status = await item.get_status(workflow=Build) if status.is_parked and status.gate == ReviewWork.name: await ReviewWork.answer( item, @@ -71,9 +71,9 @@ async def pr_review_answers_the_gate(*, repo: str, pr_number: int, payload: dict async def pr_close_settles_the_item(*, repo: str, pr_number: int, payload: dict) -> None: """GitHub announcing the verdict on a PR druks owns — one path for every merge, druks's own included. A stored verdict makes a redelivery a no-op.""" - item = WorkItem.get_for_pr(repo=repo, pr_number=pr_number, branch=payload["branch"]) + item = await WorkItem.get_for_pr(repo=repo, pr_number=pr_number, branch=payload["branch"]) if item and not item.resolution: - item.resolve(merged=payload["merged"], at=payload["resolved_at"]) + await item.resolve(merged=payload["merged"], at=payload["resolved_at"]) if payload["merged"]: await item.ship() else: @@ -83,6 +83,6 @@ async def pr_close_settles_the_item(*, repo: str, pr_number: int, payload: dict) @subscribe("ticket.transitioned") async def ticket_transition_drives_the_funnel(*, payload: dict) -> None: """Dispatch a build when a ticket from the chosen tracker enters its trigger status.""" - settings = Ship.settings() + settings = await Ship.settings() if payload["source"] == settings.tracker and payload["status"] == settings.trigger_status: await Build.dispatch(ticket=payload) diff --git a/backend/druks/contrib/ship/webhooks.py b/backend/druks/contrib/ship/webhooks.py index b1188308..59121b9b 100644 --- a/backend/druks/contrib/ship/webhooks.py +++ b/backend/druks/contrib/ship/webhooks.py @@ -18,9 +18,9 @@ class LinearEvents(Webhook): provider = "linear" category = "events" - def request_is_authentic(self) -> bool: + async def request_is_authentic(self) -> bool: try: - row = services.Linear.get() + row = await services.Linear.get() except ServiceNotConnectedError as error: raise HTTPException( status.HTTP_401_UNAUTHORIZED, @@ -94,9 +94,9 @@ class JiraEvents(Webhook): provider = "jira" category = "events" - def request_is_authentic(self) -> bool: + async def request_is_authentic(self) -> bool: try: - webhook_secret = services.Jira.get().secrets["webhook_secret"] + webhook_secret = (await services.Jira.get()).secrets["webhook_secret"] except ServiceNotConnectedError as error: raise HTTPException( status.HTTP_401_UNAUTHORIZED, @@ -138,7 +138,7 @@ async def on_issue_event(self) -> Response: "identifier": key, "status": issue_status["name"], "title": fields["summary"], - "url": self._issue_url(key), + "url": await self._issue_url(key), "project_name": fields["project"]["name"], "labels": fields["labels"], "assignee_email": assignee.get("emailAddress"), @@ -151,7 +151,7 @@ async def on_issue_event(self) -> Response: ) return JSONResponse({"accepted": True}) - def _issue_url(self, key: str) -> str: + async def _issue_url(self, key: str) -> str: # Dispatch runs after request_is_authentic, so the row is connected here. - base_url = services.Jira.get().identity["base_url"] + base_url = (await services.Jira.get()).identity["base_url"] return f"{base_url.rstrip('/')}/browse/{key}" diff --git a/backend/druks/contrib/ship/workflows.py b/backend/druks/contrib/ship/workflows.py index 06ad6e2d..dcf4df23 100644 --- a/backend/druks/contrib/ship/workflows.py +++ b/backend/druks/contrib/ship/workflows.py @@ -106,18 +106,22 @@ class Settings(BaseModel): async def dispatch(cls, *, ticket: dict) -> str | None: # The tracker funnel's entry: a ticket at the trigger status opens a build. # Resolve-or-refresh the item, then start (start() dedups a live run). - item = WorkItem.get_for_ticket_key(source=ticket["source"], ticket_key=ticket["identifier"]) + item = await WorkItem.get_for_ticket_key( + source=ticket["source"], ticket_key=ticket["identifier"] + ) if item: if item.resolution == "merged": logger.info( "Ticket %s is already merged; skipping redelivery.", ticket["identifier"] ) return - item.update(title=ticket["title"], ticket_url=ticket["url"]) + await item.update(title=ticket["title"], ticket_url=ticket["url"]) else: - repo = ProjectRepo.lookup(project_name=ticket["project_name"], labels=ticket["labels"]) + repo = await ProjectRepo.lookup( + project_name=ticket["project_name"], labels=ticket["labels"] + ) if repo: - item = WorkItem.create( + item = await WorkItem.create( project_id=repo.project_id, source=ticket["source"], title=ticket["title"] or ticket["identifier"], @@ -129,7 +133,7 @@ async def dispatch(cls, *, ticket: dict) -> str | None: logger.info("Ticket %s has no routable repo; skipping.", ticket["identifier"]) return try: - ServiceIdentity.get(GITHUB) + await ServiceIdentity.get(GITHUB) except ServiceNotConnectedError as error: # A raise would 5xx the tracker's webhook and put the delivery into # provider redelivery; the delivery itself succeeded. Log the @@ -137,7 +141,7 @@ async def dispatch(cls, *, ticket: dict) -> str | None: logger.info("Ticket %s cannot start a build: %s", ticket["identifier"], error) return email = ticket["assignee_email"] - assignee = Account.get_for_username(email.strip()) if email else None + assignee = await Account.get_for_username(email.strip()) if email else None return await cls.start( subject=item, account_id=assignee.id if assignee else None, @@ -172,11 +176,11 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: # ones they actually need under get_related_root (the prompt names them, the # credential helper handles auth). The mkdir keeps Claude's --add-dir target # valid before the first on-demand clone. - repo = self.subject.repo + repo = (await self.subject).repo # Planning agents run before the first implement provisions the branch — their # VMs clone the default branch; every agent after delivery gets the PR branch. branch = self.branch - github_token = await get_github_client().token_for_repo(repo) + github_token = await (await get_github_client()).token_for_repo(repo) await sandbox.write_secret( secret=github_token, remote=get_github_token_remote_path(sandbox.ssh_username) ) @@ -188,7 +192,7 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: ) await sandbox.exec(["mkdir", "-p", get_related_root(sandbox.ssh_username)], timeout=10.0) try: - mcp_token = await get_review_actor().client.token_for_repo(repo) + mcp_token = await (await get_review_actor()).client.token_for_repo(repo) except Exception as error: # There is no build without github: agents push and review through # the github MCP, so a run that can't mint its token fails here, @@ -207,8 +211,8 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: } async def get_prompt_context(self, **context: Any) -> dict[str, Any]: - work_item = self.subject - target_repo = ProjectRepo.get_for_repo(work_item.repo, raise_on_missing=True) + work_item = await self.subject + target_repo = await ProjectRepo.get_for_repo(work_item.repo, raise_on_missing=True) endpoint = load_settings().urls.endpoint.rstrip("/") work_item_url = f"{endpoint}/ship/work-items/{work_item.id}" if endpoint else "" prompt_context = BuildPromptContext( @@ -221,10 +225,10 @@ async def get_prompt_context(self, **context: Any) -> dict[str, Any]: issue_number=self.input.issue_number, task_owner_name=self.input.task_owner_name, task_owner_email=self.input.task_owner_email, - related_repos=target_repo.siblings(), - skills=Skill.list_delivered(self._profile.get("recommended_skills", [])), + related_repos=await target_repo.siblings(), + skills=await Skill.list_delivered(self._profile.get("recommended_skills", [])), review_code=self._settings.review_code, - review_mode=get_review_actor().mode, + review_mode=(await get_review_actor()).mode, journal=self.journal, ) return { @@ -238,9 +242,9 @@ async def get_prompt_context(self, **context: Any) -> dict[str, Any]: @step async def _load_policy_and_profile(self) -> dict[str, Any]: # One memoized read: the live policy + the work item's repo profiled facts. - repo = self.subject.repo + repo = (await self.subject).repo policy = await RepoPolicy.resolve(repo) - target = ProjectRepo.get_for_repo(repo, raise_on_missing=True) + target = await ProjectRepo.get_for_repo(repo, raise_on_missing=True) return { "policy": policy.model_dump(mode="json"), "profile": target.effective_profile, @@ -249,7 +253,7 @@ async def _load_policy_and_profile(self) -> dict[str, Any]: @step async def _load_settings(self) -> "Build.Settings": # A step so replay reuses the values the run started with, not later edits. - return self.settings() + return await self.settings() async def _plan_phase(self) -> bool: """True → implement.""" @@ -329,7 +333,7 @@ async def _approved_work(self) -> bool: return True logger.warning( "GitHub did not accept the merge of %s#%s; re-parking for review.", - self.subject.repo, + (await self.subject).repo, self.pr_number, ) return await self._work_gate() @@ -366,8 +370,8 @@ async def implement(self) -> ImplementationOutput: @step async def declare_merge_intent(self) -> bool: """Whether GitHub accepted ownership of the merge.""" - github = get_github_client() - return await github.merge_when_ready(self.subject.repo, self.pr_number) + github = await get_github_client() + return await github.merge_when_ready((await self.subject).repo, self.pr_number) # The provisioned branch + PR, pinned to the FIRST delivery — None until then # (planning runs against the default branch, and there is no PR to point at). @@ -389,10 +393,10 @@ async def _clear_draft(self) -> None: async def request_assignee_review(self) -> None: login = self.journal.assignee_github_login - repo = self.subject.repo + repo = (await self.subject).repo if login and self.pr_number: try: - await get_github_client().request_pull_request_reviewers( + await (await get_github_client()).request_pull_request_reviewers( repo, self.pr_number, [login] ) except Exception: # noqa: BLE001 — a missed ping must not fail the park @@ -404,10 +408,10 @@ async def request_assignee_review(self) -> None: ) async def set_pr_draft(self, *, draft: bool) -> None: - repo = self.subject.repo + repo = (await self.subject).repo if self.pr_number: try: - await get_github_client().set_pull_request_draft_state( + await (await get_github_client()).set_pull_request_draft_state( repo, self.pr_number, draft=draft ) except Exception: # noqa: BLE001 — a draft merge fails loudly anyway @@ -429,7 +433,7 @@ async def dispatch(cls, repo: ProjectRepo, *, refresh_only: bool = False) -> str # The profiler clones with an operator-App token, so resolve the # identity before the start spends a run and provisions a VM — the # raising lookup surfaces the actionable not-connected error. - ServiceIdentity.get(GITHUB) + await ServiceIdentity.get(GITHUB) return await cls.start( subject=repo, repo_id=repo.id, @@ -439,7 +443,7 @@ async def dispatch(cls, repo: ProjectRepo, *, refresh_only: bool = False) -> str async def run(self, repo_id: int, refresh_only: bool = False) -> None: # Every dispatch site verifies the repo exists first; a build never # profiles a repo that isn't there. - project_repo = ProjectRepo.get(repo_id) + project_repo = await ProjectRepo.get(repo_id) if refresh_only: baseline = project_repo.profile.get("baseline") or {} @@ -447,7 +451,7 @@ async def run(self, repo_id: int, refresh_only: bool = False) -> None: baseline = await Ship.repo_profiler(repo=project_repo.full_name) # The agent picks from the catalog it was handed, but a skill can be # disabled between prompt render and result — read the ground truth again. - enabled = {skill.name for skill in Skill.list_enabled()} + enabled = {skill.name for skill in await Skill.list_enabled()} baseline["recommended_skills"] = [ name for name in baseline["recommended_skills"] if name in enabled ] @@ -458,11 +462,11 @@ async def run(self, repo_id: int, refresh_only: bool = False) -> None: effective["verification"] = policy.verification.get_commands( detected=baseline.get("verification") or {} ) - project_repo.set_profile(baseline=baseline, effective=effective) + await project_repo.set_profile(baseline=baseline, effective=effective) async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: - repo = ProjectRepo.get(self.input.repo_id).full_name - github_token = await get_github_client().token_for_repo(repo) + repo = (await ProjectRepo.get(self.input.repo_id)).full_name + github_token = await (await get_github_client()).token_for_repo(repo) await sandbox.write_secret( secret=github_token, remote=get_github_token_remote_path(sandbox.ssh_username) ) @@ -480,10 +484,10 @@ async def get_workspace_kwargs(self, sandbox: "Sandbox") -> dict[str, Any]: async def get_prompt_context(self, **context: Any) -> dict[str, Any]: return { - "repo": ProjectRepo.get(self.input.repo_id).full_name, + "repo": (await ProjectRepo.get(self.input.repo_id)).full_name, "skills_catalog": [ {"name": skill.name, "description": skill.description} - for skill in Skill.list_enabled() + for skill in await Skill.list_enabled() ], **await super().get_prompt_context(**context), } diff --git a/backend/druks/core/apis/github.py b/backend/druks/core/apis/github.py index 230bc4b1..1d7eb6d1 100644 --- a/backend/druks/core/apis/github.py +++ b/backend/druks/core/apis/github.py @@ -568,13 +568,13 @@ def _fold_comments_into_body(body: str, comments: list[ReviewComment]) -> str: return "\n".join(lines) -def get_github_client() -> GitHubClient: +async def get_github_client() -> GitHubClient: """The operator client, resolved from the GitHub service-identity row — the only credential source; there is no settings or file fallback. Raises ``ServiceNotConnectedError`` when GitHub isn't connected. ``github_api_url`` stays a Settings input because it is transport, not identity. PEM plaintext exists only here, feeding the client's auth strategy.""" - row = ServiceIdentity.get(GITHUB) + row = await ServiceIdentity.get(GITHUB) return GitHubClient( app_id=row.identity["app_id"], private_key=row.secrets["private_key"], diff --git a/backend/druks/core/routes.py b/backend/druks/core/routes.py index c16061c1..af44e6fc 100644 --- a/backend/druks/core/routes.py +++ b/backend/druks/core/routes.py @@ -60,7 +60,7 @@ async def github_manifest_callback(request: Request, code: str = "") -> HTMLResp ) app = converted.json() slug = app["slug"] - ServiceIdentity.connect( + await ServiceIdentity.connect( GITHUB, identity={"app_id": str(app["id"]), "slug": slug}, secrets={"private_key": app["pem"], "webhook_secret": app["webhook_secret"]}, diff --git a/backend/druks/core/tasks.py b/backend/druks/core/tasks.py index c334b424..cae29bf8 100644 --- a/backend/druks/core/tasks.py +++ b/backend/druks/core/tasks.py @@ -20,19 +20,19 @@ async def refresh_tokens() -> None: @task(every="0 6 * * *") async def refresh_models() -> None: - fallback_id = UserSettings.get().fallback_account_id + fallback_id = (await UserSettings.get()).fallback_account_id for harness in get_harnesses(): - connections = HarnessConnection.list_for_harness(harness.name) + connections = await HarnessConnection.list_for_harness(harness.name) if not connections: continue preferred = [c for c in connections if c.account_id == fallback_id] - settings = HarnessSettings.require(harness.name) + settings = await HarnessSettings.require(harness.name) await settings.refresh_models((preferred or connections)[0]) async def _refresh() -> dict[str, object]: by_name = {harness.name: harness for harness in get_harnesses()} - connections = [c for c in HarnessConnection.list_all() if c.harness in by_name] + connections = [c for c in await HarnessConnection.list_all() if c.harness in by_name] # A refresh 401s a VM mid-call holding the old token, so a due rotation # runs only while its connection is idle — busy defers to the next tick; diff --git a/backend/druks/core/webhooks/github.py b/backend/druks/core/webhooks/github.py index aed46d9e..eb21691c 100644 --- a/backend/druks/core/webhooks/github.py +++ b/backend/druks/core/webhooks/github.py @@ -28,12 +28,12 @@ class GitHubEvents(Webhook): EVENT_HEADER: ClassVar[str] = "x-github-event" DELIVERY_HEADER: ClassVar[str] = "x-github-delivery" - def request_is_authentic(self) -> bool: + async def request_is_authentic(self) -> bool: # The delivery secret lives on the GitHub service-identity row — the # same paste that connected the App. No identity, no secret to verify # against: reject before any event dispatch. try: - identity = ServiceIdentity.get(GITHUB) + identity = await ServiceIdentity.get(GITHUB) except ServiceNotConnectedError as error: raise HTTPException( status.HTTP_401_UNAUTHORIZED, diff --git a/backend/druks/core/webhooks/slack.py b/backend/druks/core/webhooks/slack.py index 614a7f3c..ce93b320 100644 --- a/backend/druks/core/webhooks/slack.py +++ b/backend/druks/core/webhooks/slack.py @@ -70,7 +70,7 @@ class SlackInteractivity(Webhook): provider = "slack" category = "interactivity" - def request_is_authentic(self) -> bool: + async def request_is_authentic(self) -> bool: verify_slack_signature( self.raw_body, self.request.headers.get("x-slack-signature"), diff --git a/backend/druks/database.py b/backend/druks/database.py index dab82d26..35d7cae1 100644 --- a/backend/druks/database.py +++ b/backend/druks/database.py @@ -1,10 +1,15 @@ import asyncio -from collections.abc import Iterator -from contextlib import contextmanager +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from pathlib import Path from sqlalchemy import create_engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_scoped_session, + async_sessionmaker, + create_async_engine, +) _ALEMBIC_INI = Path(__file__).resolve().parent.parent / "alembic.ini" @@ -77,17 +82,24 @@ def _app_migration_dirs() -> list[tuple[str, Path]]: def create_engine_from_url(database_url: str): - # Normal transactional engine: one transaction per request/task, committed - # at the lifecycle boundary (the API session dependency, the worker session - # wrapper) so a failed unit of work rolls back instead of leaving partial - # writes. Model methods ``flush()``; the boundary commits. Low pool_timeout - # because a checkout wait blocks the event loop. The pool serves every - # concurrent run's steps plus request handling at once: a modest steady - # pool, with overflow doing the burst work — overflow connections open on - # demand and close on return, so the ceiling is high while idle cost is - # not. Ceiling 50 keeps the appliance (with DBOS's two engines at 20 each) - # inside Postgres's default 100 connections. - return create_engine( + # Synchronous engine for one-shot processes outside the event loop: the + # migrate step's seeding, the test harness's schema setup. The running app + # uses create_async_engine_from_url. + return create_engine(database_url, pool_pre_ping=True) + + +def create_async_engine_from_url(database_url: str): + # The app's engine: one transaction per request/task, committed at the + # lifecycle boundary (the API session dependency, the step session) so a + # failed unit of work rolls back instead of leaving partial writes. Model + # methods ``flush()``; the boundary commits. A checkout wait suspends the + # task, so exhaustion is backpressure — the low pool_timeout still bounds + # it. The pool serves every concurrent run's steps plus request handling + # at once: a modest steady pool, with overflow doing the burst work — + # overflow connections open on demand and close on return, so the ceiling + # is high while idle cost is not. Ceiling 50 keeps the appliance (with + # DBOS's two engines at 20 each) inside Postgres's default 100 connections. + return create_async_engine( database_url, pool_pre_ping=True, pool_timeout=5, @@ -96,8 +108,8 @@ def create_engine_from_url(database_url: str): ) -def get_session(engine) -> Session: - return Session(engine, autocommit=False, autoflush=True, expire_on_commit=False) +def get_session(engine) -> AsyncSession: + return AsyncSession(engine, autoflush=True, expire_on_commit=False) def _session_scope() -> object | None: @@ -107,30 +119,30 @@ def _session_scope() -> object | None: return None -_session_factory = sessionmaker(class_=Session, autoflush=True, expire_on_commit=False) -db_session: scoped_session = scoped_session(_session_factory, scopefunc=_session_scope) +_session_factory = async_sessionmaker(class_=AsyncSession, autoflush=True, expire_on_commit=False) +db_session: async_scoped_session = async_scoped_session(_session_factory, scopefunc=_session_scope) def configure_session(engine) -> None: _session_factory.configure(bind=engine) -@contextmanager -def session_scope(engine) -> Iterator[None]: +@asynccontextmanager +async def session_scope(engine) -> AsyncIterator[None]: """Bind a fresh DB session to the ``db_session`` registry for the block, removing it on exit — for work that runs outside the request/task session boundary (launch's schedule reconcile, a stream's per-poll snapshot), so it can't leak a session per viewer. Commits on success like the request boundary — a bare Session close rolls back, silently discarding the block's writes.""" - with get_session(engine) as session: + async with get_session(engine) as session: db_session.registry.set(session) try: yield except BaseException: - session.rollback() + await session.rollback() raise else: - session.commit() + await session.commit() finally: - db_session.remove() + await db_session.remove() diff --git a/backend/druks/doctor.py b/backend/druks/doctor.py index cde86c6e..6238527c 100644 --- a/backend/druks/doctor.py +++ b/backend/druks/doctor.py @@ -1,9 +1,11 @@ import asyncio import importlib +import inspect import os import pkgutil import socket import time +from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path @@ -20,7 +22,7 @@ from .apps.loader import iter_apps from .apps.registry import _ROLES, agents, autodiscover, services, webhooks, workflows from .core.apis.github import get_github_client -from .database import create_engine_from_url, db_session +from .database import create_async_engine_from_url, create_engine_from_url, session_scope from .harnesses.models import HarnessConnection from .harnesses.registry import get_harnesses from .sandbox.client import sandbox_client @@ -43,22 +45,32 @@ class CheckResult: pending: bool = False -def check_service_identities(settings: Settings) -> list[CheckResult]: +@asynccontextmanager +async def _check_engine(settings: Settings): + # One seam for every DB-reading check; the suite patches it to hand the + # fixture's connection in instead of a fresh engine. + engine = create_async_engine_from_url(settings.database_url) + try: + yield engine + finally: + await engine.dispose() + + +async def check_service_identities(settings: Settings) -> list[CheckResult]: """One result per declared service, so a newly-declared one is covered without editing doctor. Declarations self-register through the same discovery walk the loader runs. Identities are database-backed like harness connections — this reports each row's presence, not a file or setting.""" for app in iter_apps(): app.discover() - engine = create_engine_from_url(settings.database_url) - try: - with Session(engine) as session: - db_session.registry.set(session) + + async def _read() -> list[CheckResult]: + async with _check_engine(settings) as engine, session_scope(engine): results: list[CheckResult] = [] for service in services.all(): name = f"{service.slug}_identity" try: - row = ServiceIdentity.get(service.slug) + row = await ServiceIdentity.get(service.slug) except ServiceNotConnectedError: results.append( CheckResult( @@ -73,28 +85,31 @@ def check_service_identities(settings: Settings) -> list[CheckResult]: facts = " ".join(f"{key}={value}" for key, value in row.identity.items()) results.append(CheckResult(name=name, ok=True, detail=f"connected; {facts}")) return results + + try: + return await _read() except Exception as error: # noqa: BLE001 — a DB-read failure is a fail, not a crash return [ CheckResult( name="service_identities", ok=False, detail=f"cannot read the identities: {error}" ) ] - finally: - db_session.remove() - engine.dispose() -def check_installations(settings: Settings) -> CheckResult: +async def check_installations(settings: Settings) -> CheckResult: """Where druks may act = the operator App's installation accounts; this check is the audit surface for that set. The zero-argument client factory reads the service-identity row, so a one-off Session is bound into the ambient ``db_session`` registry for the duration.""" - engine = create_engine_from_url(settings.database_url) + + async def _list_accounts() -> tuple[str, ...]: + async with _check_engine(settings) as engine: + async with session_scope(engine): + client = await get_github_client() + return await client.list_installation_accounts() + try: - with Session(engine) as session: - db_session.registry.set(session) - client = get_github_client() - accounts = asyncio.run(client.list_installation_accounts()) + accounts = await _list_accounts() except ServiceNotConnectedError: return CheckResult( name="installations", @@ -108,9 +123,6 @@ def check_installations(settings: Settings) -> CheckResult: ok=False, detail=f"could not list operator App installations: {exc}", ) - finally: - db_session.remove() - engine.dispose() if not accounts: return CheckResult( name="installations", @@ -235,7 +247,7 @@ def check_database(settings: Settings) -> CheckResult: return CheckResult(name="database", ok=True, detail="reachable") -def check_drukbox(settings: Settings) -> CheckResult: +async def check_drukbox(settings: Settings) -> CheckResult: if not settings.sandbox.service_url: return CheckResult( name="drukbox", @@ -243,7 +255,7 @@ def check_drukbox(settings: Settings) -> CheckResult: detail="not configured (deployments: [sandbox].service_url in druks.toml)", ) try: - report = asyncio.run(_drukbox_doctor(settings)) + report = await _drukbox_doctor(settings) except Exception as error: # noqa: BLE001 — surface any SDK/transport failure as fail return CheckResult(name="drukbox", ok=False, detail=f"unreachable: {error}") if report.ok: @@ -267,7 +279,7 @@ async def _drukbox_doctor(settings: Settings): await api.aclose() -def check_sandbox_e2e(settings: Settings) -> CheckResult: +async def check_sandbox_e2e(settings: Settings) -> CheckResult: """Provision a real VM and exercise the two dial paths builds use: the acquire-time connection and a reattach from a GET-built record. Costs one VM-minute — opt-in via ``druks doctor --sandbox``, never @@ -275,7 +287,7 @@ def check_sandbox_e2e(settings: Settings) -> CheckResult: if not settings.sandbox.service_url: return CheckResult(name="sandbox_e2e", ok=True, detail="not configured") try: - detail = asyncio.run(_sandbox_e2e()) + detail = await _sandbox_e2e() except Exception as error: # noqa: BLE001 — doctor reports, never raises return CheckResult(name="sandbox_e2e", ok=False, detail=f"{error}") return CheckResult(name="sandbox_e2e", ok=True, detail=detail) @@ -398,20 +410,19 @@ def check_capability_modules(settings: Settings) -> CheckResult: ) -def check_apps(settings: Settings) -> list[CheckResult]: +async def check_apps(settings: Settings) -> list[CheckResult]: """Each installed app's resolved settings and own checks, namespaced under it. Read off the class headlessly through the loader, so doctor never imports an app's private modules. A check or settings clean that raises is contained under the app's name, and core checks remain separate ``CHECKS`` entries.""" - engine = create_engine_from_url(settings.database_url) - try: - with Session(engine) as session: - db_session.registry.set(session) + + async def _read() -> list[CheckResult]: + async with _check_engine(settings) as engine, session_scope(engine): results: list[CheckResult] = [] for app in iter_apps(): if settings_model := app.settings_model: try: - problems = app.settings().clean() + problems = (await app.settings()).clean() detail = "; ".join( f"{settings_model.model_fields[field].title or field}: {message}" for field, message in problems.items() @@ -433,20 +444,21 @@ def check_apps(settings: Settings) -> list[CheckResult]: ) ) for check in app.checks or (): - results.append(_run_app_check(app.name, check)) + results.append(await _run_app_check(app.name, check)) return results - finally: - db_session.remove() - engine.dispose() + + return await _read() -def _run_app_check(app_name: str, check) -> CheckResult: +async def _run_app_check(app_name: str, check) -> CheckResult: """One app check, its result namespaced under the app. A check that raises, or returns anything but a ``CheckResult`` (a missing ``return`` yields ``None``), becomes a failing result rather than escaping and hiding later checks.""" label = getattr(check, "__name__", repr(check)) try: outcome = check() + if inspect.iscoroutine(outcome): + outcome = await outcome if not isinstance(outcome, CheckResult): raise TypeError(f"check returned {type(outcome).__name__}, expected CheckResult") except Exception as error: # noqa: BLE001 — the check fails, never aborts @@ -473,15 +485,17 @@ def _run_app_check(app_name: str, check) -> CheckResult: ) -def run_checks(settings: Settings, *, sandbox: bool = False) -> list[CheckResult]: +async def run_checks(settings: Settings, *, sandbox: bool = False) -> list[CheckResult]: # A check yields one result, or several (check_harness_credentials fans out # over the harness registry). results: list[CheckResult] = [] for check in CHECKS: outcome = check(settings) + if inspect.iscoroutine(outcome): + outcome = await outcome results.extend(outcome if isinstance(outcome, list) else [outcome]) if sandbox: - results.append(check_sandbox_e2e(settings)) + results.append(await check_sandbox_e2e(settings)) return results @@ -520,4 +534,4 @@ def main(*, sandbox: bool = False) -> int: print() print("doctor: could not load Settings. Fix the configuration and re-run.") return 1 - return print_results(run_checks(settings, sandbox=sandbox)) + return print_results(asyncio.run(run_checks(settings, sandbox=sandbox))) diff --git a/backend/druks/durable/datastructures.py b/backend/druks/durable/datastructures.py index 2f64cfbd..8d293bbb 100644 --- a/backend/druks/durable/datastructures.py +++ b/backend/druks/durable/datastructures.py @@ -39,7 +39,7 @@ def label(self) -> str: return self.id @classmethod - def get_for_subject_id(cls, subject_id: str) -> Self | None: + async def get_for_subject_id(cls, subject_id: str) -> Self | None: """The subject this id names. Ids reach the read side as free text off a URL, so override to return None for a shape this subject could never wear.""" return cls(id=subject_id) @@ -52,7 +52,7 @@ def get_summary(self) -> "SubjectSummary": ) @classmethod - def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": + async def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": """The subjects on this class's board, newest-movement first, each as its domain summary. ``account_id`` is the caller, or None outside a request. A shared board ignores it. Returns a covariant ``Sequence`` so an app can return @@ -62,18 +62,18 @@ def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": f"a workflow declares {cls.__name__}, so it needs a list_summaries()" ) - def get_status(self, *, workflow: "type[Workflow] | None" = None) -> "SubjectStatus": + async def get_status(self, *, workflow: "type[Workflow] | None" = None) -> "SubjectStatus": """Where this subject stands: the state of the run driving it, narrowed to one workflow's runs when a subject has several kinds in flight.""" from druks.durable.reads import get_subject_status - return get_subject_status(self.subject_type, self.id, workflow=workflow) + return await get_subject_status(self.subject_type, self.id, workflow=workflow) - def get_timeline(self) -> "list[RunResponse]": + async def get_timeline(self) -> "list[RunResponse]": """Every run about this subject, oldest first, each with its agent calls.""" from druks.durable.reads import list_subject_timeline - return list_subject_timeline(self.subject_type, self.id) + return await list_subject_timeline(self.subject_type, self.id) async def get_phase(self) -> str | None: """The step it is on right now ("provisioning_vm"), while something is running.""" @@ -82,7 +82,7 @@ async def get_phase(self) -> str | None: return await get_subject_phase(self.subject_type, self.id) @classmethod - def list_open(cls, *, limit: int = 50) -> list[Self]: + async def list_open(cls, *, limit: int = 50) -> list[Self]: """The subjects of this class whose newest run hasn't handed off — still going, or failed and wanting the operator. What an app's active view lists when the subject is identity alone; a subject with rows of its own lists them with @@ -91,5 +91,5 @@ def list_open(cls, *, limit: int = 50) -> list[Self]: from druks.database import db_session from druks.durable.models import Run - open_ids = db_session().scalars(Run.open_subject_ids(cls.subject_type).limit(limit)) + open_ids = await db_session().scalars(Run.open_subject_ids(cls.subject_type).limit(limit)) return [cls(id=subject_id) for subject_id in open_ids] diff --git a/backend/druks/durable/engine.py b/backend/druks/durable/engine.py index 0d7a5905..8ca34b9d 100644 --- a/backend/druks/durable/engine.py +++ b/backend/druks/durable/engine.py @@ -5,9 +5,9 @@ from typing import TYPE_CHECKING, Any from dbos import DBOS, DBOSConfig, Queue -from sqlalchemy.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession -from druks.database import create_engine_from_url, db_session, get_session, session_scope +from druks.database import create_async_engine_from_url, db_session, get_session, session_scope from druks.durable.dbos_state import DBOS_SYSTEM_SCHEMA from druks.settings import load_settings from druks.user_settings.models import UserSettings @@ -68,7 +68,7 @@ async def _sched_entry(_scheduled_at: datetime, context: dict[str, Any] | None = _scheduled.append((cls, _sched_entry)) -def apply_schedules() -> None: +async def apply_schedules() -> None: # Declared crons name the schedule set; the operator's settings overrides only # retune or pause a declared name, never add one — so an undeclared sys-db # schedule is a renamed/removed cron: drop it. The workflow class owns its @@ -76,27 +76,29 @@ def apply_schedules() -> None: # launch() binds one, and the settings route that just wrote an override # re-runs this on its request session). declared = {cls.kind for cls, _ in _scheduled} - for existing in DBOS.list_schedules(): + for existing in await DBOS.list_schedules_async(): if existing["schedule_name"] not in declared: - DBOS.delete_schedule(existing["schedule_name"]) + await DBOS.delete_schedule_async(existing["schedule_name"]) # Crons fire on the operator's clock: "daily at midnight" means their # midnight. Evaluating in-zone (rather than converting to a UTC cron once) # keeps wall-clock cadences honest across DST. The timezone setting is # validated at its write boundary, so it's a real IANA name here. - timezone = UserSettings.get().timezone + timezone = (await UserSettings.get()).timezone for cls, fn in _scheduled: - DBOS.delete_schedule(cls.kind) - cron = cls.get_schedule() - if cls.has_enabled_schedule() and cron: - DBOS.create_schedule( + await DBOS.delete_schedule_async(cls.kind) + cron = await cls.get_schedule() + if await cls.has_enabled_schedule() and cron: + await DBOS.create_schedule_async( schedule_name=cls.kind, workflow_fn=fn, schedule=cron, cron_timezone=timezone ) -def launch() -> None: +async def launch() -> None: + # Called with the serving loop running, so DBOS captures it as the main + # loop and async steps share it. DBOS.launch() - with session_scope(_step_engine()): - apply_schedules() + async with session_scope(_step_engine()): + await apply_schedules() def shutdown() -> None: @@ -117,22 +119,22 @@ def configure_engine(engine) -> None: def _step_engine(): global _engine if not _engine: - _engine = create_engine_from_url(load_settings().database_url) + _engine = create_async_engine_from_url(load_settings().database_url) return _engine @asynccontextmanager -async def step_session() -> AsyncIterator[Session]: +async def step_session() -> AsyncIterator[AsyncSession]: # One transaction per durable step (the body itself does no IO). session = get_session(_step_engine()) db_session.registry.set(session) try: yield session except BaseException: - session.rollback() + await session.rollback() raise else: - session.commit() + await session.commit() finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() diff --git a/backend/druks/durable/models.py b/backend/druks/durable/models.py index 2bf2ffdd..7584faf2 100644 --- a/backend/druks/durable/models.py +++ b/backend/druks/durable/models.py @@ -110,25 +110,27 @@ def failure_message(self) -> str | None: return self.agent_calls[-1].last_error @classmethod - def create_row(cls, engine, *, workflow_id: str, kind: str, account_id: str | None) -> None: + async def create_row( + cls, engine, *, workflow_id: str, kind: str, account_id: str | None + ) -> None: # Own committed transaction (not the caller's request txn) so the row # exists before the running workflow's first lifecycle event. Idempotent: # a scheduled run creates its row inside the (replayable) body, and a # start that races its own retry must not double-insert. - with get_session(engine) as session: - session.execute( + async with get_session(engine) as session: + await session.execute( pg_insert(cls) .values(id=workflow_id, kind=kind, account_id=account_id or SYSTEM_ACCOUNT_ID) .on_conflict_do_nothing() ) - session.commit() + await session.commit() @classmethod - def get(cls, workflow_id: str) -> "Run | None": - return db_session().get(cls, workflow_id) + async def get(cls, workflow_id: str) -> "Run | None": + return await db_session().get(cls, workflow_id) @classmethod - def list_for_subject( + async def list_for_subject( cls, subject_type: str, subject_id: str, @@ -152,10 +154,10 @@ def list_for_subject( stmt = stmt.where(cls.kind == kind) if include_calls: stmt = stmt.options(selectinload(cls.agent_calls)) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def get_latest_for_subject( + async def get_latest_for_subject( cls, subject_type: str, subject_id: str, kind: str | None = None ) -> "Run | None": """The run that speaks for the subject: a subject holds at most one active @@ -169,7 +171,7 @@ def get_latest_for_subject( ) if kind: stmt = stmt.where(cls.kind == kind) - return db_session().scalars(stmt).first() + return (await db_session().scalars(stmt)).first() @classmethod def open_subject_ids(cls, subject_type: str) -> Select: @@ -249,7 +251,7 @@ def get_open_subjects(cls) -> Select: .order_by(driving.c.created_at.desc(), driving.c.run_id.desc()) ) - def get_ask(self) -> dict[str, Any]: + async def get_ask(self) -> dict[str, Any]: # The parked ask, ready to serve. An in-app review's ask names neither # label nor artifact — a parked run can't produce new artifacts, so the # latest resolves here on demand; fields the gate declared win. @@ -258,20 +260,20 @@ def get_ask(self) -> dict[str, Any]: raise ValueError(f"run {self.id} is not parked on an ask") if ask.get("presentation") != "in_app": return ask - artifact = Artifact.get_latest_for_run(self.id) + artifact = await Artifact.get_latest_for_run(self.id) return { "label": f"Review: {artifact.title}" if artifact else "Review", "artifact_id": artifact.id if artifact else None, **ask, } - def get_rendered_ask(self) -> dict[str, Any]: + async def get_rendered_ask(self) -> dict[str, Any]: # The parked ask as notification content. The body is the ask's own # prose (plain text, never a template); actions are data — deliver() # renders the buttons and encodes the token. Shape-dispatch on # presentation happens before any subscripting, so the external branch # never touches in-app-only keys. - ask = self.get_ask() + ask = await self.get_ask() if ask["presentation"] == "in_app": questions = ask["questions"] lines = [ask["label"], *(question["prompt"] for question in questions)] @@ -289,13 +291,13 @@ def get_rendered_ask(self) -> dict[str, Any]: # actions; url is an optional gate-author-declared view-link. return {"body": ask["label"], "actions": None, "deep_link": ask.get("url")} - def create_park_notification(self, destination_id: str, subject: dict[str, Any]) -> str: + async def create_park_notification(self, destination_id: str, subject: dict[str, Any]) -> str: # Create the notification for the round this run just parked on — the # caller supplies the run's subject and enqueues delivery. run_id + # run_parked_at snapshot the round so a click on an old button can be # refused once the run re-parks. - rendered = self.get_rendered_ask() - notification = Notification.create( + rendered = await self.get_rendered_ask() + notification = await Notification.create( destination_id=destination_id, reason="gate.parked", body=rendered["body"], @@ -326,10 +328,9 @@ async def resume(self, **fields: Any) -> None: idempotency_key=f"{self.input_gate}:{self.input_requested_at}", ) - @property - def subject(self) -> dict[str, str] | None: + async def get_subject(self) -> dict[str, str] | None: # Stamped at start; a subjectless cron has none. - attributes = db_session().scalar( + attributes = await db_session().scalar( select(workflow_status.c.attributes).where(workflow_status.c.workflow_uuid == self.id) ) if attributes: @@ -346,11 +347,11 @@ async def cancel(self, *, failure: str | None = None) -> None: self.input_gate = None self.input_request = None self.failure = failure - db_session().flush() + await db_session().flush() await DBOS.cancel_workflow_async(self.id) # The body raises DBOSWorkflowCancelledError and re-raises without # emitting, so the canceller announces the terminal state itself. - subject = self.subject + subject = await self.get_subject() if subject: await publish( WorkflowEvent.CANCELLED, @@ -371,13 +372,13 @@ async def retry(self) -> str: queue_name=run_queue.name, ) workflow_id = handle.workflow_id - Run.create_row( + await Run.create_row( _step_engine(), workflow_id=workflow_id, kind=self.kind, account_id=self.account_id, ) - subject = self.subject + subject = await self.get_subject() if subject: await publish( WorkflowEvent.RETRIED, @@ -496,7 +497,7 @@ def get_stream_path(self, stream: Literal["stdout", "stderr"]) -> Path | None: return candidate if candidate.exists() else None @classmethod - def start( + async def start( cls, engine, *, @@ -512,11 +513,11 @@ def start( # shows while the agent works — the running step's session won't commit # until it ends. Provisioning isn't part of the call, so the row only # exists once there's a host to run on. - with get_session(engine) as session: + async with get_session(engine) as session: # A crash-recovered step re-runs with a fresh call id; abandon the # prior attempt's RUNNING row (a run's calls are sequential) so it # doesn't linger as a phantom live step. - session.execute( + await session.execute( update(cls) .where(cls.run_id == run_id, cls.status == AgentCallStatus.RUNNING.value) .values(status=AgentCallStatus.ABANDONED.value, finished_at=Base.utc_now()) @@ -531,12 +532,12 @@ def start( account_id=account_id, ) ) - session.commit() + await session.commit() @classmethod - def finish(cls, engine, *, call_id: str, result: "AgentResult") -> None: - with get_session(engine) as session: - call = session.get(cls, call_id) + async def finish(cls, engine, *, call_id: str, result: "AgentResult") -> None: + async with get_session(engine) as session: + call = await session.get(cls, call_id) call.status = result.status.value call.started_at = result.started_at call.finished_at = Base.utc_now() @@ -544,43 +545,43 @@ def finish(cls, engine, *, call_id: str, result: "AgentResult") -> None: call.failure_code = result.error.code if result.error else "" call.cost_usd = result.cost_usd call.cost_metadata = result.cost_metadata - session.commit() + await session.commit() @classmethod - def fail(cls, engine, *, call_id: str, error: BaseException) -> None: + async def fail(cls, engine, *, call_id: str, error: BaseException) -> None: # The run raised after the call started (a cancel, or a crash past the # agent body) — close the row so it doesn't linger as a phantom step. - with get_session(engine) as session: - call = session.get(cls, call_id) + async with get_session(engine) as session: + call = await session.get(cls, call_id) call.status = AgentCallStatus.FAILED.value call.finished_at = Base.utc_now() call.last_error = str(error) - session.commit() + await session.commit() @classmethod - def get(cls, agent_call_id: str) -> "AgentCall": - call = db_session().get(cls, agent_call_id) + async def get(cls, agent_call_id: str) -> "AgentCall": + call = await db_session().get(cls, agent_call_id) if not call: raise AgentCallNotFound(agent_call_id) return call @classmethod - def list_for_run(cls, run_id: str) -> list["AgentCall"]: + async def list_for_run(cls, run_id: str) -> list["AgentCall"]: # Execution order — the same order Run.agent_calls loads. stmt = select(cls).where(cls.run_id == run_id).order_by(cls.created_at, cls.id) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def list_for_subject(cls, subject_type: str, subject_id: str) -> list["AgentCall"]: + async def list_for_subject(cls, subject_type: str, subject_id: str) -> list["AgentCall"]: stmt = ( select(cls) .where(subject_filter(cls.run_id, subject_type, subject_id)) .order_by(cls.created_at, cls.id) ) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def total_run_spend_between(cls, *, start: datetime, end: datetime) -> tuple[float, int]: + async def total_run_spend_between(cls, *, start: datetime, end: datetime) -> tuple[float, int]: stmt = ( select(cls.cost_usd, cls.cost_metadata) .where(cls.finished_at.is_not(None)) @@ -589,7 +590,7 @@ def total_run_spend_between(cls, *, start: datetime, end: datetime) -> tuple[flo ) cost = 0.0 tokens = 0 - for cost_usd, metadata in db_session().execute(stmt): + for cost_usd, metadata in await db_session().execute(stmt): if cost_usd is not None: cost += float(cost_usd) canonical = normalize_token_usage(metadata) @@ -597,14 +598,14 @@ def total_run_spend_between(cls, *, start: datetime, end: datetime) -> tuple[flo tokens += canonical["total_tokens"] return cost, tokens - def record_cost(self, *, cost_usd: float | None, cost_metadata: dict | None) -> None: + async def record_cost(self, *, cost_usd: float | None, cost_metadata: dict | None) -> None: if cost_usd is None and not cost_metadata: return if cost_usd is not None: self.cost_usd = cost_usd if cost_metadata: self.cost_metadata = cost_metadata - db_session().flush() + await db_session().flush() class Artifact(Base, Uuid7Pk): @@ -621,7 +622,9 @@ class Artifact(Base, Uuid7Pk): path: Mapped[str] @classmethod - def record(cls, *, call_dir: Path, call_id: str, kind: str, title: str, content: str) -> None: + async def record( + cls, *, call_dir: Path, call_id: str, kind: str, title: str, content: str + ) -> None: # Platform-owned: write a call's declared renderable output into its dir and # record the descriptor on the call's step session. Idempotent per call # (unique fk) so a replayed step never double-records. @@ -629,22 +632,22 @@ def record(cls, *, call_dir: Path, call_id: str, kind: str, title: str, content: call_dir.mkdir(parents=True, exist_ok=True) (call_dir / name).write_text(content) session = db_session() - session.execute( + await session.execute( pg_insert(cls) .values(agent_call_id=call_id, kind=kind, title=title, path=name) .on_conflict_do_nothing(index_elements=["agent_call_id"]) ) - session.flush() + await session.flush() @classmethod - def get_for_call(cls, call_id: str) -> "Artifact | None": - return db_session().scalar(select(cls).where(cls.agent_call_id == call_id)) + async def get_for_call(cls, call_id: str) -> "Artifact | None": + return await db_session().scalar(select(cls).where(cls.agent_call_id == call_id)) @classmethod - def get_latest_for_run(cls, run_id: str) -> "Artifact | None": + async def get_latest_for_run(cls, run_id: str) -> "Artifact | None": # The run's most recent renderable output, reached through its calls — an # in-app review shows this beside its controls. Newest call wins. - return db_session().scalar( + return await db_session().scalar( select(cls) .join(AgentCall, AgentCall.id == cls.agent_call_id) .where(AgentCall.run_id == run_id) diff --git a/backend/druks/durable/reads.py b/backend/druks/durable/reads.py index f04f4768..456a2c02 100644 --- a/backend/druks/durable/reads.py +++ b/backend/druks/durable/reads.py @@ -37,9 +37,9 @@ _TERMINAL_CALL_STATES = {"succeeded", "failed", "abandoned"} -def get_agent_call_files(call_id: str) -> AgentCallFiles: - call = AgentCall.get(call_id) - artifact = Artifact.get_for_call(call.id) +async def get_agent_call_files(call_id: str) -> AgentCallFiles: + call = await AgentCall.get(call_id) + artifact = await Artifact.get_for_call(call.id) layout = call.artifact_layout def named(path: Path) -> ArtifactFile | None: @@ -67,30 +67,30 @@ def named(path: Path) -> ArtifactFile | None: ) -def list_subject_timeline(subject_type: str, subject_id: str) -> list[RunResponse]: +async def list_subject_timeline(subject_type: str, subject_id: str) -> list[RunResponse]: # The subject's whole timeline: every run about it, oldest first, each # with its agent calls. - runs = Run.list_for_subject(subject_type, subject_id, include_calls=True) - return _timeline(runs) + runs = await Run.list_for_subject(subject_type, subject_id, include_calls=True) + return await _timeline(runs) -def get_subject_status( +async def get_subject_status( subject_type: str, subject_id: str, *, workflow: "type[Workflow] | None" = None ) -> SubjectStatus: kind = workflow.kind if workflow else None - latest = Run.get_latest_for_subject(subject_type, subject_id, kind=kind) - return _status(latest) + latest = await Run.get_latest_for_subject(subject_type, subject_id, kind=kind) + return await _status(latest) async def get_subject_phase(subject_type: str, subject_id: str) -> str | None: - runs = Run.list_for_subject(subject_type, subject_id) + runs = await Run.list_for_subject(subject_type, subject_id) active_run = next((run for run in runs if run.is_active), None) if active_run and active_run.is_running: return await get_run_phase(active_run.id) return -def get_subject_response( +async def get_subject_response( subject_type: str, subject_id: str, *, @@ -100,28 +100,28 @@ def get_subject_response( # list_for_subject is newest-first, so runs[0] is the driving run the status # reads — the same row get_latest_for_subject would return, its calls already # eager-loaded here. - runs = Run.list_for_subject(subject_type, subject_id, include_calls=True) + runs = await Run.list_for_subject(subject_type, subject_id, include_calls=True) latest = runs[0] if runs else None return SubjectResponse( summary=summary, - status=_status(latest), - timeline=_timeline(runs), + status=await _status(latest), + timeline=await _timeline(runs), activity=activity, ) -def _timeline(runs: list[Run]) -> list[RunResponse]: +async def _timeline(runs: list[Run]) -> list[RunResponse]: ordered = sorted(runs, key=lambda run: (run.created_at, run.id)) return [ RunResponse.from_run( run, - input_request=run.get_ask() if run.input_request else None, + input_request=await run.get_ask() if run.input_request else None, ) for run in ordered ] -def _status(driving_run: Run | None) -> SubjectStatus: +async def _status(driving_run: Run | None) -> SubjectStatus: # Facts only: the app's UI renders its copy from them. if not driving_run: return SubjectStatus(state=RunState.SCHEDULED) @@ -134,8 +134,10 @@ def _status(driving_run: Run | None) -> SubjectStatus: # inverse: only a parked run's input_gate is a live ask (a timed-out run # keeps the stale column). agent = None - if not parked and driving_run.agent_calls: - agent = driving_run.agent_calls[-1].agent + if not parked: + calls = await driving_run.awaitable_attrs.agent_calls + if calls: + agent = calls[-1].agent return SubjectStatus( state=RunState(driving_run.state), kind=driving_run.kind, @@ -214,9 +216,9 @@ async def stream_transcript( elapsed = 0.0 last_keepalive = 0.0 while True: - with session_scope(engine): + async with session_scope(engine): try: - call = AgentCall.get(call_id) + call = await AgentCall.get(call_id) except AgentCallNotFound: return path = call.get_stream_path(stream) diff --git a/backend/druks/events/builder.py b/backend/druks/events/builder.py index 8f936156..e31c44e2 100644 --- a/backend/druks/events/builder.py +++ b/backend/druks/events/builder.py @@ -8,20 +8,20 @@ _FETCH_LIMIT = 500 -def build_feed( +async def build_feed( *, app: str | None = None, before: int | None = None, limit: int = _PAGE_LIMIT_DEFAULT, ) -> tuple[list[FeedItem], str | None]: - items = [FeedItem.model_validate(event) for event in _events(app, before)] + items = [FeedItem.model_validate(event) for event in await _events(app, before)] items.sort(key=lambda item: item.seq, reverse=True) page = items[:limit] next_cursor = str(page[-1].seq) if len(page) == limit and page else None return page, next_cursor -def _events(app: str | None, before: int | None) -> list[Event]: +async def _events(app: str | None, before: int | None) -> list[Event]: # This app's events plus any unscoped (core) ones. The log stores the app; # the core never derives it from the subject. stmt = select(Event).order_by(Event.id.desc()) @@ -29,4 +29,4 @@ def _events(app: str | None, before: int | None) -> list[Event]: stmt = stmt.where(Event.id < before) if app: stmt = stmt.where(or_(Event.app == app, Event.app.is_(None))) - return list(db_session().scalars(stmt.limit(_FETCH_LIMIT)).all()) + return list((await db_session().scalars(stmt.limit(_FETCH_LIMIT))).all()) diff --git a/backend/druks/events/models.py b/backend/druks/events/models.py index 1b90d7b9..bf6baff6 100644 --- a/backend/druks/events/models.py +++ b/backend/druks/events/models.py @@ -34,7 +34,7 @@ class Event(Base): payload: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict) @classmethod - def emit( + async def emit( cls, *, type: str, @@ -54,4 +54,4 @@ def emit( payload=payload or {}, ) ) - db_session().flush() + await db_session().flush() diff --git a/backend/druks/events/routes.py b/backend/druks/events/routes.py index da3ab944..c2288d85 100644 --- a/backend/druks/events/routes.py +++ b/backend/druks/events/routes.py @@ -38,7 +38,7 @@ async def list_feed( app: str | None = Query(default=None), ) -> FeedResponse: cursor = _parse_cursor(before) - items, next_cursor = build_feed(app=app, before=cursor, limit=limit) + items, next_cursor = await build_feed(app=app, before=cursor, limit=limit) return FeedResponse(items=items, next_cursor=next_cursor) @@ -56,8 +56,8 @@ async def feed_stream(): return # New Session per tick so we don't hold a transaction open across the # sleep; the open/close cost is irrelevant against the poll cadence. - with session_scope(engine): - items, _next_cursor = build_feed( + async with session_scope(engine): + items, _next_cursor = await build_feed( app=app, before=None, limit=100 if first else 50, diff --git a/backend/druks/harnesses/base.py b/backend/druks/harnesses/base.py index b198be63..ea1d86ff 100644 --- a/backend/druks/harnesses/base.py +++ b/backend/druks/harnesses/base.py @@ -102,7 +102,7 @@ def __init__( self.sandbox = sandbox @abstractmethod - def build_invocation(self, **kwargs: object) -> AgentInvocation: + async def build_invocation(self, **kwargs: object) -> AgentInvocation: """Assemble this CLI's full invocation (argv, stdin, credentials, env) for one prompt. Pure — never touches the live sandbox; the sandbox executes the returned invocation.""" @@ -123,7 +123,7 @@ def check_returncode(cls, result: HarnessRunResult) -> None: raise error(message) raise exceptions.HarnessError(message) - def get_manifest( + async def get_manifest( self, *, mcp_servers: tuple["McpServer", ...], @@ -150,7 +150,7 @@ def get_manifest( # token_present reads the delivered env: a server's bearer env var is # set iff its token was found at delivery, for a static or an # app-minted token alike. - declared = {server["name"]: server for server in mcp_models.McpServer.list_enabled()} + declared = {server["name"]: server for server in await mcp_models.McpServer.list_enabled()} delivered_by_name = {server.name: server for server in mcp_servers} mcp = [] for name in sorted(declared.keys() | delivered_by_name.keys()): @@ -173,7 +173,7 @@ def get_manifest( "model": self.model or "", "harness": self.name, "mcp_servers": mcp, - "skills_delivered": sorted(skill.name for skill in Skill.list_delivered(skills)), + "skills_delivered": sorted(skill.name for skill in await Skill.list_delivered(skills)), } canonical = json.dumps(capability, sort_keys=True, separators=(",", ":")) return { @@ -196,10 +196,10 @@ def mint_run_id(call_id: str | None) -> str: return call_id or str(uuid.uuid4()) @classmethod - def get_credentials(cls) -> dict: + async def get_credentials(cls) -> dict: """The fallback account's credential dict — for callers with no selection.""" - row = HarnessConnection.get_for_account(cls.name, fallback=True) + row = await HarnessConnection.get_for_account(cls.name, fallback=True) data = dict(row.payload) if row else None if data: return data @@ -208,12 +208,12 @@ def get_credentials(cls) -> dict: ) @classmethod - def render_credentials_file(cls, connection_id: str | None = None) -> str: + async def render_credentials_file(cls, connection_id: str | None = None) -> str: """The selected connection's payload, read fresh at push time; a vanished row fails the call rather than render another account's.""" if not connection_id: - return json.dumps(cls.get_credentials()) - row = HarnessConnection.get(connection_id) + return json.dumps(await cls.get_credentials()) + row = await HarnessConnection.get(connection_id) if row: return json.dumps(dict(row.payload)) raise exceptions.HarnessNotConnectedError( @@ -319,7 +319,7 @@ async def rotate_token( reports ``locked`` without touching the provider — two concurrent grants on one refresh lineage trip the provider's reuse detection.""" moment = now or _utc_now() - row = HarnessConnection.reload(connection_id) + row = await HarnessConnection.reload(connection_id) if not row: return RotationResult( cls.name, "failed", error="no_credentials", connection_id=connection_id @@ -342,7 +342,7 @@ async def rotate_token( try: # Re-read after winning the lock: the previous holder may have # advanced this lineage (or deleted the row) after our first read. - row = HarnessConnection.reload(connection_id) + row = await HarnessConnection.reload(connection_id) if not row: return RotationResult( cls.name, "failed", error="no_credentials", connection_id=connection_id @@ -365,8 +365,8 @@ async def rotate_token( # presenting it again can never succeed. Drop only this # credential so the connection reads as disconnected — the # UI shows Reconnect and the next tick has no row to hammer. - row.delete() - db_session().commit() + await row.delete() + await db_session().commit() logger.warning( "%s connection %s auto-disconnected after invalid_grant; " "reconnect to restore", @@ -379,13 +379,13 @@ async def rotate_token( cls.name, "failed", error="bad_response", connection_id=row.id ) - row.update_payload(data, expires_at=new_expiry) + await row.update_payload(data, expires_at=new_expiry) # The grant is externally anchored — the provider may have killed # the old refresh token the moment it issued this one — so the new # lineage must be committed before the lock releases; deferring to # the step's own commit would let a concurrent refresher take the # freed lock and re-present the superseded token. - db_session().commit() + await db_session().commit() return RotationResult( cls.name, "refreshed", expires_at=new_expiry, connection_id=row.id ) @@ -474,7 +474,7 @@ async def poll_usage(cls, connection: HarnessConnection) -> dict[str, object]: parsed = await cls.fetch_usage(connection) except Exception: # noqa: BLE001 — a crashed scrape records an error row, not a failed refresh logger.warning("usage fetch crashed for %s", cls.name, exc_info=True) - UsageScrape( + await UsageScrape( harness=cls.name, account_id=account_id, parse_ok=False, @@ -502,7 +502,7 @@ async def poll_usage(cls, connection: HarnessConnection) -> dict[str, object]: snapshot.five_hour_percent_left = parsed.five_hour.percent_left snapshot.five_hour_resets_at = parsed.five_hour.resets_at snapshot.weeks = _WEEKLY_WINDOWS.dump_python(parsed.weeks, mode="json") - snapshot.save() + await snapshot.save() return { "harness": cls.name, "account_id": account_id, diff --git a/backend/druks/harnesses/claude.py b/backend/druks/harnesses/claude.py index 9b8d53a9..b2bd8025 100644 --- a/backend/druks/harnesses/claude.py +++ b/backend/druks/harnesses/claude.py @@ -72,7 +72,7 @@ class ClaudeHarness(Harness): "api error: 5": exceptions.HarnessOverloadedError, } - def build_invocation( + async def build_invocation( self, *, prompt: str, @@ -140,7 +140,7 @@ def build_invocation( name="claude", args=("sh", "-c", wrapper), stdin=prompt.encode("utf-8"), - credentials=_claude_credentials( + credentials=await _claude_credentials( self.sandbox, github_token=github_token, include_plugins=include_plugins, @@ -404,7 +404,7 @@ def _parse_iso(value: object) -> datetime | None: return ensure_utc(parsed) -def _claude_credentials( +async def _claude_credentials( sandbox: SandboxSettings, *, github_token: str | None, @@ -455,11 +455,11 @@ def _claude_credentials( if skills_src: dirs += ((skills_src, ".claude/skills"),) return Credentials( - claude_credentials=ClaudeHarness.render_credentials_file(connection_id), + claude_credentials=await ClaudeHarness.render_credentials_file(connection_id), github_token=github_token, extra_config_files=files, extra_config_dirs=dirs, - extra_dir_excludes={".claude/skills": Skill.delivery_excludes(skills)}, + extra_dir_excludes={".claude/skills": await Skill.delivery_excludes(skills)}, ) diff --git a/backend/druks/harnesses/codex.py b/backend/druks/harnesses/codex.py index 5856238c..d05d6d34 100644 --- a/backend/druks/harnesses/codex.py +++ b/backend/druks/harnesses/codex.py @@ -565,7 +565,7 @@ def _build_codex_wrapper( return ["sh", "-c", wrapper] - def build_invocation( + async def build_invocation( self, *, prompt: str, @@ -606,7 +606,7 @@ def build_invocation( name=self.name, args=tuple(cmd), stdin=_with_final_message_note(prompt).encode("utf-8"), - credentials=self._codex_credentials( + credentials=await self._codex_credentials( github_token=github_token, skills=skills, connection_id=connection_id, @@ -684,7 +684,7 @@ def _prompt_flags(self) -> tuple[str, ...]: args = (*args, "--json") return args - def _codex_credentials( + async def _codex_credentials( self, *, github_token: str | None, @@ -710,11 +710,11 @@ def _codex_credentials( if skills_src: dirs = ((skills_src, ".codex/skills"),) return Credentials( - codex_credentials=self.render_credentials_file(connection_id), + codex_credentials=await self.render_credentials_file(connection_id), github_token=github_token, extra_config_files=files, extra_config_dirs=dirs, - extra_dir_excludes={".codex/skills": Skill.delivery_excludes(skills)}, + extra_dir_excludes={".codex/skills": await Skill.delivery_excludes(skills)}, ) diff --git a/backend/druks/harnesses/models.py b/backend/druks/harnesses/models.py index 9cf3d319..e1673e66 100644 --- a/backend/druks/harnesses/models.py +++ b/backend/druks/harnesses/models.py @@ -32,19 +32,19 @@ class HarnessConnection(Base, Uuid7Pk): updated_at: Mapped[datetime] = mapped_column(default=Base.utc_now, onupdate=Base.utc_now) @classmethod - def get(cls, connection_id: str) -> "HarnessConnection | None": - return db_session().get(cls, connection_id) + async def get(cls, connection_id: str) -> "HarnessConnection | None": + return await db_session().get(cls, connection_id) @classmethod - def lookup(cls, harness: str, account_id: str | None) -> "HarnessConnection": + async def lookup(cls, harness: str, account_id: str | None) -> "HarnessConnection": """The connection a call runs with: the account's own, else the fallback account's — which carries unmatched work so automation keeps moving.""" if account_id: - own = cls.get_for_account(harness, account_id) + own = await cls.get_for_account(harness, account_id) if own: return own - fallback = cls.get_for_account(harness, fallback=True) + fallback = await cls.get_for_account(harness, fallback=True) if fallback: return fallback raise HarnessNotConnectedError( @@ -53,42 +53,42 @@ def lookup(cls, harness: str, account_id: str | None) -> "HarnessConnection": ) @classmethod - def get_for_account( + async def get_for_account( cls, harness: str, account_id: str | None = None, *, fallback: bool = False ) -> "HarnessConnection | None": """``fallback=True`` resolves the fallback account's connection — what actor-less execution runs as.""" if fallback: - account_id = UserSettings.get().fallback_account_id - return db_session().scalar( + account_id = (await UserSettings.get()).fallback_account_id + return await db_session().scalar( select(cls).where(cls.harness == harness, cls.account_id == account_id) ) @classmethod - def list_all(cls) -> list["HarnessConnection"]: - return list(db_session().scalars(select(cls).order_by(cls.harness, cls.id))) + async def list_all(cls) -> list["HarnessConnection"]: + return list(await db_session().scalars(select(cls).order_by(cls.harness, cls.id))) @classmethod - def list_for_account(cls, account_id: str) -> list["HarnessConnection"]: + async def list_for_account(cls, account_id: str) -> list["HarnessConnection"]: stmt = select(cls).where(cls.account_id == account_id).order_by(cls.harness) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def list_for_harness(cls, harness: str) -> list["HarnessConnection"]: + async def list_for_harness(cls, harness: str) -> list["HarnessConnection"]: stmt = select(cls).where(cls.harness == harness).order_by(cls.id) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def reload(cls, connection_id: str) -> "HarnessConnection | None": + async def reload(cls, connection_id: str) -> "HarnessConnection | None": """Fresh-from-DB read of one row, past the identity map's cached state — the post-lock re-read that keeps a refresher from re-presenting a refresh token a concurrent winner already advanced.""" - return db_session().scalar( + return await db_session().scalar( select(cls).where(cls.id == connection_id).execution_options(populate_existing=True) ) @classmethod - def connect( + async def connect( cls, *, harness: str, @@ -100,17 +100,17 @@ def connect( """Upsert ``account``'s connection for this harness — update its existing row or create one.""" session = db_session() - row = cls.get_for_account(harness, account.id) + row = await cls.get_for_account(harness, account.id) if not row: row = cls(harness=harness, account_id=account.id) session.add(row) row.payload = payload row.provider_email = provider_email row.expires_at = expires_at - session.flush() + await session.flush() return row - def update_payload(self, payload: dict, *, expires_at: datetime | None) -> None: + async def update_payload(self, payload: dict, *, expires_at: datetime | None) -> None: # Whole-value reassignment is the write path: the encrypted column # re-encrypts what it's handed. The caller's dict may alias the live # mapping's nested blocks (a dict() copy is shallow), making old and @@ -120,9 +120,9 @@ def update_payload(self, payload: dict, *, expires_at: datetime | None) -> None: self.payload = payload flag_modified(self, "payload") self.expires_at = expires_at - db_session().flush() + await db_session().flush() - def delete(self) -> None: + async def delete(self) -> None: session = db_session() - session.delete(self) - session.flush() + await session.delete(self) + await session.flush() diff --git a/backend/druks/harnesses/registry.py b/backend/druks/harnesses/registry.py index 9b96c7b1..8a35dd1a 100644 --- a/backend/druks/harnesses/registry.py +++ b/backend/druks/harnesses/registry.py @@ -16,7 +16,7 @@ def get_harness(name: str) -> type[Harness] | None: return harness -def get_harness_for_model(model: str) -> type[Harness]: +async def get_harness_for_model(model: str) -> type[Harness]: """The harness that runs ``model`` from its fetched-or-shipped model list. A miss raises loudly; namespace-shaped models do not route until a fetch @@ -24,7 +24,7 @@ def get_harness_for_model(model: str) -> type[Harness]: """ from druks.user_settings.models import HarnessSettings - for row in HarnessSettings.all(): + for row in await HarnessSettings.all(): if model == row.name or any(m["id"] == model for m in row.allowed_models): return row.harness raise HarnessError(f"no harness runs model {model!r}") diff --git a/backend/druks/harnesses/routes.py b/backend/druks/harnesses/routes.py index 07db83d3..d031471f 100644 --- a/backend/druks/harnesses/routes.py +++ b/backend/druks/harnesses/routes.py @@ -67,13 +67,13 @@ async def complete_connection( # provider-verified email; get_or_create is atomic, so concurrent # completions of the same email converge, and a true different-email # race surfaces as the none-mode multi-operator refusal. - resolved = account or Account.get_or_create(completed.provider_email) + resolved = account or await Account.get_or_create(completed.provider_email) # Runs with no actor execute as the fallback account; claim the slot when # none is set yet. - settings = UserSettings.get() + settings = await UserSettings.get() if not settings.fallback_account_id: - settings.set_fallback_account(resolved.id) - connection = HarnessConnection.connect( + await settings.set_fallback_account(resolved.id) + connection = await HarnessConnection.connect( harness=harness.name, account=resolved, payload=completed.payload, @@ -85,16 +85,16 @@ async def complete_connection( # writer on this event loop, and nothing past the point of durability may # depend on another database read. response = AccountResponse.model_validate(resolved) - db_session().commit() + await db_session().commit() try: # Fresh picker right after connect; fetch failures are tagged inside. # The single-use flow is already spent, so trouble here — including a # database that vanished under the refresh — only logs. - await HarnessSettings.require(harness.name).refresh_models(connection) + await (await HarnessSettings.require(harness.name)).refresh_models(connection) except Exception: logging.getLogger(__name__).exception("Model refresh after connect failed") with suppress(Exception): - db_session().rollback() + await db_session().rollback() return response @@ -103,8 +103,8 @@ async def disconnect_harness( name: str, account: Account = Depends(current_session_account) ) -> HarnessResponse: harness = _resolve_harness(name) - connection = HarnessConnection.get_for_account(harness.name, account.id) + connection = await HarnessConnection.get_for_account(harness.name, account.id) if connection: # Only the requesting account's own connection — never another's. - connection.delete() - return HarnessResponse.from_row(HarnessSettings.require(harness.name), None, account) + await connection.delete() + return HarnessResponse.from_row(await HarnessSettings.require(harness.name), None, account) diff --git a/backend/druks/mcp/gateway/routes.py b/backend/druks/mcp/gateway/routes.py index 3598f00b..c354c4ac 100644 --- a/backend/druks/mcp/gateway/routes.py +++ b/backend/druks/mcp/gateway/routes.py @@ -29,7 +29,7 @@ async def get_gate( ) -> schemas.GateResponse: """A parked run's open gate: the ask, a bounded artifact chunk, and parkedAt — echo parkedAt unchanged to answer_gate.""" - return services.get_gate(run) + return await services.get_gate(run) @router.post( @@ -75,7 +75,7 @@ async def get_agent_call( ) -> schemas.AgentCallDetailResponse: """One agent call's metadata with bounded transcript and stderr tails and an artifact chunk.""" - return services.get_agent_call(call) + return await services.get_agent_call(call) @router.get( @@ -87,4 +87,4 @@ async def get_agent_call( async def get_usage(account: Account = Depends(current_account)) -> schemas.AgentUsageResponse: """The caller's harness quota snapshot and today's spend. Pure read — it never triggers a scrape.""" - return services.get_usage(account) + return await services.get_usage(account) diff --git a/backend/druks/mcp/gateway/services.py b/backend/druks/mcp/gateway/services.py index 54186ce9..9e9bb23e 100644 --- a/backend/druks/mcp/gateway/services.py +++ b/backend/druks/mcp/gateway/services.py @@ -27,8 +27,8 @@ _HISTORY_POINTS = 8 -def get_gate(run_id: str) -> schemas.GateResponse: - run = Run.get(run_id) +async def get_gate(run_id: str) -> schemas.GateResponse: + run = await Run.get(run_id) if not run: raise RunNotFound(run_id) if run.state != RunState.PARKED.value: @@ -40,18 +40,18 @@ def get_gate(run_id: str) -> schemas.GateResponse: run=run.id, gate=run.input_gate, # type: ignore[arg-type] parked_at=run.input_requested_at, # type: ignore[arg-type] - ask=run.get_ask(), - artifact=_artifact_content(Artifact.get_latest_for_run(run.id)), + ask=await run.get_ask(), + artifact=await _artifact_content(await Artifact.get_latest_for_run(run.id)), ) async def answer_gate( run_id: str, *, parked_at: datetime, control: str, answers: dict[str, str], note: str ) -> schemas.GateAnswerResponse: - run = Run.get(run_id) + run = await Run.get(run_id) if not run: raise RunNotFound(run_id) - db_session().expire(run) # the receipt/park comparison must read fresh + await db_session().refresh(run) # the receipt/park comparison must read fresh if run.answer_parked_at == parked_at: return schemas.GateAnswerResponse( run=run.id, parked_at=parked_at, result="already_answered" @@ -64,15 +64,15 @@ async def answer_gate( if not ask or ask.get("presentation") != "in_app": raise exceptions.GateNotAnswerable(run_id) try: - payload = validate_in_app_answer(run.get_ask(), control, answers, note) + payload = validate_in_app_answer(await run.get_ask(), control, answers, note) except InvalidChoiceError as error: raise exceptions.InvalidGateAnswer(str(error)) from error await run.resume(**payload) return schemas.GateAnswerResponse(run=run.id, parked_at=parked_at, result="answered") -def get_agent_call(call_id: str) -> schemas.AgentCallDetailResponse: - call = AgentCall.get(call_id) +async def get_agent_call(call_id: str) -> schemas.AgentCallDetailResponse: + call = await AgentCall.get(call_id) layout = call.artifact_layout return schemas.AgentCallDetailResponse( run=call.run_id, @@ -81,15 +81,15 @@ def get_agent_call(call_id: str) -> schemas.AgentCallDetailResponse: layout.transcript, offset=-_TRANSCRIPT_TAIL_BYTES, limit=_TRANSCRIPT_TAIL_BYTES ).text, stderr=read_slice(layout.stderr, offset=-_STDERR_TAIL_BYTES, limit=_STDERR_TAIL_BYTES).text, - artifact=_artifact_content(Artifact.get_for_call(call.id)), + artifact=await _artifact_content(await Artifact.get_for_call(call.id)), ) -def _artifact_content(artifact: Artifact | None) -> schemas.ArtifactContent | None: +async def _artifact_content(artifact: Artifact | None) -> schemas.ArtifactContent | None: if not artifact: return try: - call = AgentCall.get(artifact.agent_call_id) + call = await AgentCall.get(artifact.agent_call_id) except AgentCallNotFound: return path = call.get_file_path(artifact.path) @@ -103,10 +103,12 @@ def _artifact_content(artifact: Artifact | None) -> schemas.ArtifactContent | No ) -def get_usage(account: Account) -> schemas.AgentUsageResponse: +async def get_usage(account: Account) -> schemas.AgentUsageResponse: now = datetime.now(UTC) - timezone, local_start = operator_local_day(UserSettings.get().timezone, now) - rows = list_finished_calls(account.id, since=local_start, until=local_start + timedelta(days=1)) + timezone, local_start = operator_local_day((await UserSettings.get()).timezone, now) + rows = await list_finished_calls( + account.id, since=local_start, until=local_start + timedelta(days=1) + ) spend = 0.0 tokens = 0 for _, cost_usd, cost_metadata, _ in rows: @@ -121,16 +123,16 @@ def get_usage(account: Account) -> schemas.AgentUsageResponse: spend_today_usd=round(spend, 4), tokens_today=tokens, runs_today=len(rows), - harnesses=[_harness_usage(h.name, account.id, now=now) for h in get_harnesses()], + harnesses=[await _harness_usage(h.name, account.id, now=now) for h in get_harnesses()], ) -def _harness_usage(name: str, account_id: str, *, now: datetime) -> schemas.AgentHarnessUsage: - is_connected = bool(HarnessConnection.get_for_account(name, account_id)) - row = UsageScrape.latest_for(name, account_id) +async def _harness_usage(name: str, account_id: str, *, now: datetime) -> schemas.AgentHarnessUsage: + is_connected = bool(await HarnessConnection.get_for_account(name, account_id)) + row = await UsageScrape.latest_for(name, account_id) if not row: return schemas.AgentHarnessUsage(name=name, is_connected=is_connected) - history = UsageScrape.history_for(name, account_id, since=now - WEEK_RANGE) + history = await UsageScrape.history_for(name, account_id, since=now - WEEK_RANGE) five_hour_cutoff = now - FIVE_HOUR_RANGE five_hour = [ UsageHistoryPoint(t=point.scraped_at, pct=point.five_hour_percent_left) diff --git a/backend/druks/mcp/models.py b/backend/druks/mcp/models.py index fc68e1ef..c32d67ac 100644 --- a/backend/druks/mcp/models.py +++ b/backend/druks/mcp/models.py @@ -46,21 +46,23 @@ class McpServer(Base, Uuid7Pk): created_at: Mapped[datetime] = mapped_column(default=Base.utc_now) @classmethod - def list_all(cls) -> list["McpServer"]: + async def list_all(cls) -> list["McpServer"]: # The raw overlay rows — not the merged registry view (_merged). - return list(db_session().execute(select(cls).order_by(cls.name)).scalars()) + return list((await db_session().execute(select(cls).order_by(cls.name))).scalars()) @classmethod - def get_for_name(cls, name: str) -> "McpServer | None": - return db_session().execute(select(cls).where(cls.name == name)).scalar_one_or_none() + async def get_for_name(cls, name: str) -> "McpServer | None": + return ( + await db_session().execute(select(cls).where(cls.name == name)) + ).scalar_one_or_none() @classmethod - def _merged(cls) -> dict[str, dict]: + async def _merged(cls) -> dict[str, dict]: # The full view the API and delivery build from, keyed by # name: each built-in definition (url + auth from the registry) # overlaid with its operator row's enable choice and secrets, then any # fully custom rows. - rows = {server.name: server for server in cls.list_all()} + rows = {server.name: server for server in await cls.list_all()} servers: dict[str, dict] = {} for definition in mcp_servers.all(): row = rows.pop(definition["name"], None) @@ -92,8 +94,8 @@ def _merged(cls) -> dict[str, dict]: return servers @classmethod - def get_resolved(cls, account_id: str | None) -> dict[str, dict]: - servers = cls._merged() + async def get_resolved(cls, account_id: str | None) -> dict[str, dict]: + servers = await cls._merged() # has_token = nothing blocks this server's auth at delivery, read from # wherever its source keeps the secret: druks' env for an env-sourced # server, a stored grant for a connected one, the stored token for a @@ -109,7 +111,7 @@ def get_resolved(cls, account_id: str | None) -> dict[str, dict]: if server["identity_mode"]: grant_account = get_grant_account(server["identity_mode"], account_id) server["has_token"] = bool( - OauthConnection.list_for_account( + await OauthConnection.list_for_account( grant_provider(server["name"]), grant_account ) ) @@ -118,26 +120,26 @@ def get_resolved(cls, account_id: str | None) -> dict[str, dict]: return servers @classmethod - def list_enabled(cls) -> list[dict]: + async def list_enabled(cls) -> list[dict]: # The enabled subset — what a run delivers and the settings UI shows active. - return [server for server in cls._merged().values() if server["is_enabled"]] + return [server for server in (await cls._merged()).values() if server["is_enabled"]] @classmethod - def set_enabled(cls, name: str, is_enabled: bool) -> bool: + async def set_enabled(cls, name: str, is_enabled: bool) -> bool: # A built-in has no row until an operator changes its state; the enable # choice creates one, carrying the built-in's url. False means the name # is neither a row nor a catalog entry. - server = cls.get_for_name(name) + server = await cls.get_for_name(name) if server: server.is_enabled = is_enabled return True if name in mcp_servers: - cls.create(name=name, url=mcp_servers.get(name)["url"], is_enabled=is_enabled) + await cls.create(name=name, url=mcp_servers.get(name)["url"], is_enabled=is_enabled) return True return False @classmethod - def create( + async def create( cls, *, name: str, @@ -161,13 +163,13 @@ def create( is_enabled=is_enabled, ) session.add(server) - session.flush() + await session.flush() return server - def delete(self) -> None: + async def delete(self) -> None: session = db_session() - session.delete(self) - session.flush() + await session.delete(self) + await session.flush() class McpClientRegistration(Base, Uuid7Pk): @@ -188,19 +190,19 @@ class McpClientRegistration(Base, Uuid7Pk): client_secret = EncryptedTextField(default="") @classmethod - def get_for_account(cls, server_name: str, account_id: str) -> "McpClientRegistration | None": + async def get_for_account( + cls, server_name: str, account_id: str + ) -> "McpClientRegistration | None": return ( - db_session() - .execute( + await db_session().execute( select(cls) .join(McpServer, McpServer.id == cls.server_id) .where(McpServer.name == server_name, cls.account_id == account_id) ) - .scalar_one_or_none() - ) + ).scalar_one_or_none() @classmethod - def store( + async def store( cls, *, server_id: str, @@ -225,9 +227,11 @@ def store( "client_secret": statement.excluded.client_secret, }, ).returning(cls) - return session.scalars(statement, execution_options={"populate_existing": True}).one() + return ( + await session.scalars(statement, execution_options={"populate_existing": True}) + ).one() - def delete(self) -> None: + async def delete(self) -> None: session = db_session() - session.delete(self) - session.flush() + await session.delete(self) + await session.flush() diff --git a/backend/druks/mcp/oauth.py b/backend/druks/mcp/oauth.py index f5b72780..78d3b24b 100644 --- a/backend/druks/mcp/oauth.py +++ b/backend/druks/mcp/oauth.py @@ -25,15 +25,15 @@ def _http() -> httpx.AsyncClient: return httpx.AsyncClient(timeout=30.0, follow_redirects=True) -def get_connection(name: str, account_id: str) -> OauthConnection | None: +async def get_connection(name: str, account_id: str) -> OauthConnection | None: # One live connection per (server, account) — MCP's policy over the # shared table. Revoked rows stay behind as history. - rows = OauthConnection.list_for_account(grant_provider(name), account_id) + rows = await OauthConnection.list_for_account(grant_provider(name), account_id) return rows[0] if rows else None -def list_connections(name: str) -> list[OauthConnection]: - return OauthConnection.list_for_provider(grant_provider(name)) +async def list_connections(name: str) -> list[OauthConnection]: + return await OauthConnection.list_for_provider(grant_provider(name)) def _origin(url: str) -> str: @@ -229,7 +229,7 @@ async def complete_connect(*, state: str, code: str) -> str: # fill the mode if unclaimed. A concurrent claim wins the row lock; the # select reads whichever choice landed, and the grant goes under it. session = db_session() - session.execute( + await session.execute( pg_insert(McpServer) .values( name=name, @@ -238,14 +238,14 @@ async def complete_connect(*, state: str, code: str) -> str: ) .on_conflict_do_nothing(index_elements=["name"]) ) - session.execute( + await session.execute( update(McpServer) .where(McpServer.name == name, McpServer.identity_mode.is_(None)) .values(identity_mode=pending["identity_mode"]) ) - server = session.scalars(select(McpServer).where(McpServer.name == name)).one() + server = (await session.scalars(select(McpServer).where(McpServer.name == name))).one() account_id = get_grant_account(server.identity_mode, pending["account_id"]) - McpClientRegistration.store( + await McpClientRegistration.store( server_id=server.id, account_id=account_id, token_endpoint=pending["token_endpoint"], @@ -255,13 +255,15 @@ async def complete_connect(*, state: str, code: str) -> str: identity = {} if pending["userinfo_endpoint"]: identity = await fetch_identity(pending["userinfo_endpoint"], tokens["access_token"]) - connection = get_connection(name, account_id) + connection = await get_connection(name, account_id) if connection: - connection.reconnect(refresh_token=tokens["refresh_token"], scopes=[], identity=identity) + await connection.reconnect( + refresh_token=tokens["refresh_token"], scopes=[], identity=identity + ) # A reconsent's stale cached token must not serve until its TTL runs out. await evict_access_token(name, account_id) else: - OauthConnection.create( + await OauthConnection.create( provider=grant_provider(name), account_id=account_id, refresh_token=tokens["refresh_token"], @@ -272,27 +274,27 @@ async def complete_connect(*, state: str, code: str) -> str: async def evict_access_token(name: str, account_id: str) -> None: - connection = get_connection(name, account_id) + connection = await get_connection(name, account_id) if connection: await OauthClient(provider=grant_provider(name)).evict_access_token(connection.id) async def disconnect(name: str, account_id: str, *, reason: str = "user") -> None: - connection = get_connection(name, account_id) + connection = await get_connection(name, account_id) if connection: await OauthClient(provider=grant_provider(name)).disconnect(connection, reason=reason) - registration = McpClientRegistration.get_for_account(name, account_id) + registration = await McpClientRegistration.get_for_account(name, account_id) if registration: - registration.delete() + await registration.delete() async def get_access_token(name: str, account_id: str) -> str: """The delivery-side token for a connected server, served by the shared engine from this server's grant — delivery never ships a server the agent can't authenticate to.""" - connection = get_connection(name, account_id) - registration = McpClientRegistration.get_for_account(name, account_id) - server = McpServer.get_for_name(name) + connection = await get_connection(name, account_id) + registration = await McpClientRegistration.get_for_account(name, account_id) + server = await McpServer.get_for_name(name) if not connection or not registration or not server: raise MissingGrantError(name, account_id) client = OauthClient( diff --git a/backend/druks/mcp/routes.py b/backend/druks/mcp/routes.py index 083f62eb..09cafd17 100644 --- a/backend/druks/mcp/routes.py +++ b/backend/druks/mcp/routes.py @@ -27,15 +27,16 @@ router = APIRouter(prefix="/api/mcp-servers", tags=["mcp-servers"]) -def _response(name: str) -> McpServerResponse: - return McpServerResponse.model_validate(McpServer.get_resolved(current_account_id.get())[name]) +async def _response(name: str) -> McpServerResponse: + resolved = await McpServer.get_resolved(current_account_id.get()) + return McpServerResponse.model_validate(resolved[name]) @router.get("", response_model=list[McpServerResponse]) async def list_mcp_servers() -> list[McpServerResponse]: return [ McpServerResponse.model_validate(server) - for server in McpServer.get_resolved(current_account_id.get()).values() + for server in (await McpServer.get_resolved(current_account_id.get())).values() ] @@ -59,7 +60,7 @@ async def add_mcp_server(body: CreateMcpServerRequest) -> McpServerResponse: status_code=409, detail=f"MCP server {body.name!r} is built-in; configure it instead of adding it.", ) - if McpServer.get_for_name(body.name): + if await McpServer.get_for_name(body.name): raise HTTPException( status_code=409, detail=f"MCP server {body.name!r} already exists; remove it first." ) @@ -73,10 +74,10 @@ async def add_mcp_server(body: CreateMcpServerRequest) -> McpServerResponse: status_code=422, detail=f"MCP server {body.name!r} needs a bearer token." ) try: - McpServer.create(name=body.name, url=body.url, token=body.token) + await McpServer.create(name=body.name, url=body.url, token=body.token) except InvalidServerNameError as error: raise HTTPException(status_code=422, detail=str(error)) from error - return _response(body.name) + return await _response(body.name) @router.post("/registry", response_model=McpServerResponse) @@ -86,7 +87,7 @@ async def install_mcp_server(body: InstallMcpServerRequest, request: Request) -> status_code=409, detail=f"MCP server {body.name!r} is built-in; configure it instead of adding it.", ) - if McpServer.get_for_name(body.name): + if await McpServer.get_for_name(body.name): raise HTTPException( status_code=409, detail=f"MCP server {body.name!r} already exists; remove it first." ) @@ -130,7 +131,7 @@ async def install_mcp_server(body: InstallMcpServerRequest, request: Request) -> token_source = TokenSource.OAUTH is_enabled = False try: - McpServer.create( + await McpServer.create( name=body.name, url=candidate["url"], token_source=token_source, @@ -140,16 +141,16 @@ async def install_mcp_server(body: InstallMcpServerRequest, request: Request) -> ) except InvalidServerNameError as error: raise HTTPException(status_code=422, detail=str(error)) from error - return _response(body.name) + return await _response(body.name) @router.patch("/{name}", response_model=McpServerResponse) async def set_mcp_server_enabled( name: str, is_enabled: bool = Body(embed=True) ) -> McpServerResponse: - if not McpServer.set_enabled(name, is_enabled): + if not await McpServer.set_enabled(name, is_enabled): raise HTTPException(status_code=404, detail=f"MCP server {name!r} not found") - return _response(name) + return await _response(name) @router.delete("/{name}", status_code=204) @@ -160,13 +161,13 @@ async def remove_mcp_server(name: str) -> None: raise HTTPException( status_code=409, detail=f"MCP server {name!r} is managed by druks; disable it instead." ) - server = McpServer.get_for_name(name) + server = await McpServer.get_for_name(name) if not server: raise HTTPException(status_code=404, detail=f"MCP server {name!r} not found") # Revoke before the server row goes — the registration lookup needs it. - for connection in oauth.list_connections(name): + for connection in await oauth.list_connections(name): await oauth.disconnect(name, connection.account_id, reason="server_removed") - server.delete() + await server.delete() @router.post("/{name}/connect", response_model=ConnectMcpServerResponse) @@ -175,10 +176,10 @@ async def connect_mcp_server( request: Request, identity_mode: Annotated[IdentityMode, Body(embed=True)], ) -> ConnectMcpServerResponse: - server = McpServer.get_resolved(current_account_id.get()).get(name) + server = (await McpServer.get_resolved(current_account_id.get())).get(name) if not server or server["token_source"] != TokenSource.OAUTH: raise HTTPException(status_code=404, detail=f"MCP server {name!r} is not an OAuth server.") - if oauth.list_connections(name) and server["identity_mode"] != identity_mode: + if await oauth.list_connections(name) and server["identity_mode"] != identity_mode: raise HTTPException( status_code=409, detail=f"MCP server {name!r} already uses {server['identity_mode']!r} identity.", @@ -222,7 +223,7 @@ async def oauth_callback(state: str = "", code: str = "", error: str = "") -> HT raise HTTPException(status_code=400, detail=str(exchange_error)) from exchange_error # Connecting is the operator's explicit "use this server" — a # connected-but-disabled server is a dead end nobody asks for. - McpServer.set_enabled(name, is_enabled=True) + await McpServer.set_enabled(name, is_enabled=True) # druks opened this tab via window.open, so the page may close itself; the # broadcast tells the settings modal to refetch before the tab goes. The # text stays for browsers that refuse the close. @@ -231,24 +232,24 @@ async def oauth_callback(state: str = "", code: str = "", error: str = "") -> HT @router.delete("/{name}/grant", status_code=204) async def disconnect_mcp_server(name: str) -> None: - server = McpServer.get_resolved(current_account_id.get()).get(name) + server = (await McpServer.get_resolved(current_account_id.get())).get(name) if not server or server["token_source"] != TokenSource.OAUTH: raise HTTPException(status_code=404, detail=f"MCP server {name!r} is not an OAuth server.") if not server["identity_mode"]: raise HTTPException(status_code=404, detail=f"MCP server {name!r} has no grant.") account_id = get_grant_account(server["identity_mode"], current_account_id.get()) - connection = oauth.get_connection(name, account_id) + connection = await oauth.get_connection(name, account_id) if not connection: raise HTTPException( status_code=404, detail=f"MCP server {name!r} has no grant for account {account_id!r}.", ) await oauth.disconnect(name, account_id) - if not oauth.list_connections(name): + if not await oauth.list_connections(name): # The last grant leaving reopens the mode choice: the next connect is # a first connect again. - server_row = McpServer.get_for_name(name) + server_row = await McpServer.get_for_name(name) if server_row: server_row.identity_mode = None if server["identity_mode"] == IdentityMode.SHARED: - McpServer.set_enabled(name, is_enabled=False) + await McpServer.set_enabled(name, is_enabled=False) diff --git a/backend/druks/mcp/server.py b/backend/druks/mcp/server.py index 9bb37d40..05e6e435 100644 --- a/backend/druks/mcp/server.py +++ b/backend/druks/mcp/server.py @@ -43,20 +43,20 @@ async def verify_token(self, token: str) -> AccessToken | None: # Auth middleware runs outside the request session boundary, so this # owns one — authenticate stamps last_used_at. try: - pat = PersonalAccessToken.authenticate(token) + pat = await PersonalAccessToken.authenticate(token) access = AccessToken( token=token, client_id=pat.token_prefix, scopes=[], claims={"account_id": pat.account_id, "pat_id": pat.id}, ) - db_session().commit() + await db_session().commit() return access except InvalidPatError: - db_session().rollback() + await db_session().rollback() return finally: - db_session.remove() + await db_session.remove() class CallerPat(httpx.Auth): diff --git a/backend/druks/models.py b/backend/druks/models.py index 8b1b196d..6c754824 100644 --- a/backend/druks/models.py +++ b/backend/druks/models.py @@ -75,7 +75,7 @@ def label(self) -> str: return self.get_label() @classmethod - def get_for_subject_id(cls, subject_id: str) -> Self | None: + async def get_for_subject_id(cls, subject_id: str) -> Self | None: """The row this subject id names. A subject id is free text and reaches the read-side straight off a URL, so an id this table could never hold is a miss rather than an error.""" @@ -85,7 +85,7 @@ def get_for_subject_id(cls, subject_id: str) -> Self | None: key = int(subject_id) except ValueError: return - return db_session().get(cls, key) + return await db_session().get(cls, key) def get_summary(self) -> "SubjectSummary": """The header its board and page show it under — the app's own fields; @@ -95,7 +95,7 @@ def get_summary(self) -> "SubjectSummary": ) @classmethod - def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": + async def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": """The rows on this class's board, newest-movement first, each as its domain summary. ``account_id`` is the caller, or None outside a request. A shared board ignores it. Returns a covariant ``Sequence`` so an app can return @@ -105,15 +105,15 @@ def list_summaries(cls, account_id: str | None) -> "Sequence[SubjectSummary]": f"a workflow declares {cls.__name__}, so it needs a list_summaries()" ) - def get_status(self, *, workflow: "type[Workflow] | None" = None) -> "SubjectStatus": + async def get_status(self, *, workflow: "type[Workflow] | None" = None) -> "SubjectStatus": from druks.durable.reads import get_subject_status - return get_subject_status(self.subject_type, str(self.id), workflow=workflow) + return await get_subject_status(self.subject_type, str(self.id), workflow=workflow) - def get_timeline(self) -> "list[RunResponse]": + async def get_timeline(self) -> "list[RunResponse]": from druks.durable.reads import list_subject_timeline - return list_subject_timeline(self.subject_type, str(self.id)) + return await list_subject_timeline(self.subject_type, str(self.id)) async def get_phase(self) -> str | None: from druks.durable.reads import get_subject_phase @@ -121,7 +121,7 @@ async def get_phase(self) -> str | None: return await get_subject_phase(self.subject_type, str(self.id)) @classmethod - def list_open(cls, *, limit: int = 50) -> list[Self]: + async def list_open(cls, *, limit: int = 50) -> list[Self]: """The rows whose newest run hasn't handed off — still going, or failed and wanting the operator. What an app's active view lists.""" # Cycle: the durable read side is built on this module's Base. @@ -137,4 +137,4 @@ def list_open(cls, *, limit: int = 50) -> list[Self]: .order_by(cls.id.desc()) .limit(limit) ) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) diff --git a/backend/druks/notifications/models.py b/backend/druks/notifications/models.py index 00b2f173..2e2c2194 100644 --- a/backend/druks/notifications/models.py +++ b/backend/druks/notifications/models.py @@ -37,19 +37,21 @@ class Destination(Base, Uuid7Pk): updated_at: Mapped[datetime] = mapped_column(default=Base.utc_now) @classmethod - def list_all(cls) -> list["Destination"]: - return list(db_session().execute(select(cls).order_by(cls.name)).scalars()) + async def list_all(cls) -> list["Destination"]: + return list((await db_session().execute(select(cls).order_by(cls.name))).scalars()) @classmethod - def get(cls, destination_id: str) -> "Destination | None": - return db_session().get(cls, destination_id) + async def get(cls, destination_id: str) -> "Destination | None": + return await db_session().get(cls, destination_id) @classmethod - def get_for_name(cls, name: str) -> "Destination | None": - return db_session().execute(select(cls).where(cls.name == name)).scalar_one_or_none() + async def get_for_name(cls, name: str) -> "Destination | None": + return ( + await db_session().execute(select(cls).where(cls.name == name)) + ).scalar_one_or_none() @classmethod - def create(cls, *, name: str, kind: str, url: str) -> "Destination": + async def create(cls, *, name: str, kind: str, url: str) -> "Destination": try: DestinationKind(kind) except ValueError as error: @@ -57,13 +59,13 @@ def create(cls, *, name: str, kind: str, url: str) -> "Destination": session = db_session() destination = cls(name=name, kind=kind, url=url) session.add(destination) - session.flush() + await session.flush() return destination - def delete(self) -> None: + async def delete(self) -> None: session = db_session() - session.delete(self) - session.flush() + await session.delete(self) + await session.flush() class Notification(Base, Uuid7Pk): @@ -103,7 +105,7 @@ class Notification(Base, Uuid7Pk): updated_at: Mapped[datetime] = mapped_column(default=Base.utc_now) @classmethod - def create( + async def create( cls, *, destination_id: str, @@ -130,48 +132,46 @@ def create( ) session = db_session() session.add(notification) - session.flush() + await session.flush() return notification @classmethod - def get(cls, notification_id: str) -> "Notification | None": - return db_session().get(cls, notification_id) + async def get(cls, notification_id: str) -> "Notification | None": + return await db_session().get(cls, notification_id) @classmethod - def list_recent(cls, limit: int = 50) -> list["Notification"]: + async def list_recent(cls, limit: int = 50) -> list["Notification"]: # uuid7 ids are time-ordered, breaking created_at ties toward the newest. stmt = select(cls).order_by(cls.created_at.desc(), cls.id.desc()).limit(limit) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def get_for_token(cls, token: str) -> "Notification | None": + async def get_for_token(cls, token: str) -> "Notification | None": return ( - db_session() - .execute(select(cls).where(cls.correlation_token == token)) - .scalar_one_or_none() - ) + await db_session().execute(select(cls).where(cls.correlation_token == token)) + ).scalar_one_or_none() @property def is_acknowledged(self) -> bool: return self.state == NotificationState.ACKNOWLEDGED - def mark_delivered(self) -> None: + async def mark_delivered(self) -> None: self.state = NotificationState.DELIVERED.value self.delivered_at = Base.utc_now() self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def mark_failed(self, reason: str) -> None: + async def mark_failed(self, reason: str) -> None: self.state = NotificationState.FAILED.value self.last_error = reason self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def mark_acknowledged(self) -> bool: + async def mark_acknowledged(self) -> bool: # Atomic claim: exactly one concurrent responder wins the transition # (the loser's duplicate send already collapsed on the DBOS round key). session = db_session() - claimed = session.execute( + claimed = await session.execute( update(Notification) .where( Notification.id == self.id, @@ -179,5 +179,5 @@ def mark_acknowledged(self) -> bool: ) .values(state=NotificationState.ACKNOWLEDGED.value, updated_at=Base.utc_now()) ) - session.expire(self) + await session.refresh(self) return claimed.rowcount == 1 diff --git a/backend/druks/notifications/outbox.py b/backend/druks/notifications/outbox.py index 9d6f69b7..03860a7d 100644 --- a/backend/druks/notifications/outbox.py +++ b/backend/druks/notifications/outbox.py @@ -30,16 +30,17 @@ def _sanitized(error: Exception) -> str: async def send_notification(notification_id: str) -> None: async def _send() -> None: async with step_session() as session: - notification = Notification.get(notification_id) + notification = await Notification.get(notification_id) # A duplicate enqueue or the replay of a completed send finds the # row delivered and stops. At-least-once: a crash after the send but # before this step checkpoints re-posts once on recovery — the click # side is idempotent, so a duplicate message is the accepted cost. if notification.state == NotificationState.DELIVERED: return - destination = Destination.get(notification.destination_id) - # deliver() runs after this session commits; expunge keeps the - # loaded destination readable past the commit's attribute expiry. + destination = await Destination.get(notification.destination_id) + # deliver() runs after this session closes; expunge keeps the + # loaded destination readable as a detached row — an attribute + # touched past the close can't lazy-load under the async session. session.expunge(destination) body = notification.body actions = notification.actions @@ -53,7 +54,7 @@ async def _send() -> None: idempotency_key=notification_id, ) async with step_session(): - Notification.get(notification_id).mark_delivered() + await (await Notification.get(notification_id)).mark_delivered() try: await DBOS.run_step_async( @@ -64,7 +65,7 @@ async def _send() -> None: async def _mark_failed() -> None: async with step_session(): - Notification.get(notification_id).mark_failed(reason) + await (await Notification.get(notification_id)).mark_failed(reason) # Terminal: record the failure and return normally — re-raising would # put the workflow into perpetual DBOS recovery for a dead endpoint. diff --git a/backend/druks/notifications/routes.py b/backend/druks/notifications/routes.py index d7749517..9251f6da 100644 --- a/backend/druks/notifications/routes.py +++ b/backend/druks/notifications/routes.py @@ -27,14 +27,14 @@ @router.get("/destinations", response_model=list[DestinationResponse]) async def list_destinations() -> list[Destination]: - return Destination.list_all() + return await Destination.list_all() @router.post("/destinations", response_model=DestinationResponse) async def create_destination(body: CreateDestinationRequest) -> Destination: if not body.name.strip(): raise HTTPException(status_code=422, detail="Destination needs a name.") - if Destination.get_for_name(body.name): + if await Destination.get_for_name(body.name): raise HTTPException( status_code=409, detail=f"Destination {body.name!r} already exists; remove it first." ) @@ -45,39 +45,39 @@ async def create_destination(body: CreateDestinationRequest) -> Destination: raise HTTPException( status_code=422, detail="URL is not a recognized notification destination." ) - return Destination.create(name=body.name, kind=body.kind.value, url=url) + return await Destination.create(name=body.name, kind=body.kind.value, url=url) @router.patch("/destinations/{destination_id}", response_model=DestinationResponse) async def set_destination_enabled( destination_id: str, is_enabled: Annotated[bool, Body(embed=True)] ) -> Destination: - destination = Destination.get(destination_id) + destination = await Destination.get(destination_id) if not destination: raise HTTPException(status_code=404, detail=f"Destination {destination_id!r} not found") destination.is_enabled = is_enabled - db_session().flush() + await db_session().flush() return destination @router.delete("/destinations/{destination_id}", status_code=204) async def delete_destination(destination_id: str) -> None: - destination = Destination.get(destination_id) + destination = await Destination.get(destination_id) if not destination: raise HTTPException(status_code=404, detail=f"Destination {destination_id!r} not found") - destination.delete() + await destination.delete() @router.get("", response_model=list[NotificationResponse]) async def list_notifications(limit: int = Query(50, ge=1, le=500)) -> list[Notification]: - return Notification.list_recent(limit) + return await Notification.list_recent(limit) # Declared after the /destinations routes: declaration order is match order, # so the id match can't swallow them. @router.get("/{notification_id}", response_model=NotificationResponse) async def get_notification(notification_id: str) -> Notification: - notification = Notification.get(notification_id) + notification = await Notification.get(notification_id) if not notification: raise HTTPException(status_code=404, detail=f"Notification {notification_id!r} not found") return notification diff --git a/backend/druks/notifications/services.py b/backend/druks/notifications/services.py index 053766bb..1a631f00 100644 --- a/backend/druks/notifications/services.py +++ b/backend/druks/notifications/services.py @@ -46,7 +46,7 @@ def validate_in_app_answer( async def respond_to_notification(token: str, choice: dict[str, Any]) -> None: - notification = Notification.get_for_token(token) + notification = await Notification.get_for_token(token) if not notification: raise UnknownTokenError() if notification.is_acknowledged: @@ -54,15 +54,15 @@ async def respond_to_notification(token: str, choice: dict[str, Any]) -> None: if not notification.run_id: # A run-less notification routes no reply. raise InvalidChoiceError("this notification does not take an answer") - run = Run.get(notification.run_id) + run = await Run.get(notification.run_id) if not run: raise CorruptCorrelationError(notification.id, notification.run_id) # The notification snapshots the round it was sent for; the answer must - # land on the run's live round — expire so the comparison reads fresh. - db_session().expire(run) + # land on the run's live round — refresh so the comparison reads fresh. + await db_session().refresh(run) if run.state != RunState.PARKED.value or run.input_requested_at != notification.run_parked_at: raise StaleRoundError() - ask = run.get_ask() + ask = await run.get_ask() if ask.get("presentation") != "in_app": # External gates are answered on their source (PR review, ticket # comment) via the existing webhook paths, never through this rail — @@ -73,7 +73,7 @@ async def respond_to_notification(token: str, choice: dict[str, Any]) -> None: ask, choice["control"], choice.get("answers", {}), choice.get("note", "") ) await run.resume(**resume_payload) - if not notification.mark_acknowledged(): + if not await notification.mark_acknowledged(): # A concurrent responder won the claim; this send already collapsed on # the DBOS round key. raise AlreadyAcknowledgedError() diff --git a/backend/druks/sandbox/host.py b/backend/druks/sandbox/host.py index 7f268942..ff3283a7 100644 --- a/backend/druks/sandbox/host.py +++ b/backend/druks/sandbox/host.py @@ -220,8 +220,8 @@ async def run_agent( settings = load_settings() # Effort/timeout fall back to the model's harness defaults. - harness_class = get_harness_for_model(model) - harness_settings = HarnessSettings.require(harness_class.name) + harness_class = await get_harness_for_model(model) + harness_settings = await HarnessSettings.require(harness_class.name) effort = effort or harness_settings.effort timeout = timeout if timeout is not None else harness_settings.timeout harness = harness_class( @@ -303,14 +303,14 @@ async def run_prompt( persist_manifest( artifact_dir, call_id=run_id, - manifest=harness.get_manifest( + manifest=await harness.get_manifest( mcp_servers=mcp_servers, skills=skills, extra_env=extra_env, ), ) - invocation = harness.build_invocation( + invocation = await harness.build_invocation( prompt=prompt, schema=schema, run_id=run_id, diff --git a/backend/druks/services/base.py b/backend/druks/services/base.py index 4ce613df..ce847656 100644 --- a/backend/druks/services/base.py +++ b/backend/druks/services/base.py @@ -46,9 +46,8 @@ def connected_at(self): return self.row.connected_at async def get_access_token(self, scopes: tuple[str, ...] = (), cached: bool = True) -> str: - return await self.service.get_oauth_client().get_access_token( - connection=self.row, scopes=scopes, cached=cached - ) + client = await self.service.get_oauth_client() + return await client.get_access_token(connection=self.row, scopes=scopes, cached=cached) async def disconnect(self) -> None: await OauthClient(provider=self.service.slug).disconnect(self.row, reason="user") @@ -72,14 +71,14 @@ def __set_name__(self, owner: type, name: str) -> None: def label(self) -> str: return f"{self.owner.name}.{self.name}" - def list_for_account(self, account_id: str) -> list[Connection]: + async def list_for_account(self, account_id: str) -> list[Connection]: return [ Connection(self.service, row) - for row in OauthConnection.list_for_account(self.service.slug, account_id) + for row in await OauthConnection.list_for_account(self.service.slug, account_id) ] - def get(self, connection_id: str) -> Connection | None: - row = OauthConnection.get(connection_id) + async def get(self, connection_id: str) -> Connection | None: + row = await OauthConnection.get(connection_id) if row and row.provider == self.service.slug and not row.revoked_at: return Connection(self.service, row) @@ -198,8 +197,8 @@ def connect_fields(cls) -> list[dict[str, Any]]: ] @classmethod - def get(cls) -> ServiceIdentity: - return ServiceIdentity.get(cls.slug) + async def get(cls) -> ServiceIdentity: + return await ServiceIdentity.get(cls.slug) @classmethod def with_scopes(cls, *scopes: str) -> ScopedService: @@ -234,13 +233,13 @@ async def get_identity(cls, access_token: str) -> dict[str, Any]: return {} @classmethod - def get_oauth_client(cls) -> OauthClient: + async def get_oauth_client(cls) -> OauthClient: """The connected identity as a configured ``OauthClient``, keyed by the service slug. Raises ``ServiceNotConnectedError`` until the operator connects the service.""" if not cls.token_endpoint: raise TypeError(f"{cls.__name__} declares no OAuth endpoints") - connected = cls.get() + connected = await cls.get() return OauthClient( provider=cls.slug, authorization_endpoint=cls.authorization_endpoint, @@ -252,9 +251,9 @@ def get_oauth_client(cls) -> OauthClient: ) @classmethod - def is_connected(cls) -> bool: + async def is_connected(cls) -> bool: try: - ServiceIdentity.get(cls.slug) + await ServiceIdentity.get(cls.slug) except ServiceNotConnectedError: return False return True @@ -288,7 +287,7 @@ async def connect(cls, payload: dict[str, Any]) -> ServiceIdentity: } if all(str(value).strip() for value in (*identity.values(), *secrets.values())): proven = await cls.verify(settings) - return ServiceIdentity.connect( + return await ServiceIdentity.connect( cls.slug, identity={**identity, **proven}, secrets=secrets ) raise ServiceConnectError("Every field is required.") diff --git a/backend/druks/services/models.py b/backend/druks/services/models.py index 2569e3e6..575efde2 100644 --- a/backend/druks/services/models.py +++ b/backend/druks/services/models.py @@ -25,25 +25,25 @@ class ServiceIdentity(Base): connected_at: Mapped[datetime] @classmethod - def get(cls, service: str) -> "ServiceIdentity": - if identity := db_session().get(cls, service): + async def get(cls, service: str) -> "ServiceIdentity": + if identity := await db_session().get(cls, service): return identity raise ServiceNotConnectedError(service) @classmethod - def connect( + async def connect( cls, service: str, *, identity: dict[str, Any], secrets: dict[str, str] ) -> "ServiceIdentity": # The caller verifies the credentials against the service first; this # trusts what it is given and overwrites whatever was connected. - row = db_session().get(cls, service) + row = await db_session().get(cls, service) if not row: row = cls(service=service) db_session().add(row) row.identity = identity row.secrets = secrets row.connected_at = Base.utc_now() - db_session().flush() + await db_session().flush() return row @@ -77,11 +77,11 @@ class OauthConnection(Base, Uuid7Pk): revoked_reason: Mapped[str] = mapped_column(default="") @classmethod - def get(cls, connection_id: str) -> "OauthConnection | None": - return db_session().get(cls, connection_id) + async def get(cls, connection_id: str) -> "OauthConnection | None": + return await db_session().get(cls, connection_id) @classmethod - def create( + async def create( cls, *, provider: str, @@ -98,13 +98,13 @@ def create( identity=identity or {}, ) db_session().add(connection) - db_session().flush() + await db_session().flush() return connection @classmethod - def list_for_account(cls, provider: str, account_id: str) -> "list[OauthConnection]": + async def list_for_account(cls, provider: str, account_id: str) -> "list[OauthConnection]": return list( - db_session().scalars( + await db_session().scalars( select(cls) .where( cls.provider == provider, @@ -116,13 +116,12 @@ def list_for_account(cls, provider: str, account_id: str) -> "list[OauthConnecti ) @classmethod - def get_for_identity( + async def get_for_identity( cls, provider: str, account_id: str, key: str, value: Any ) -> "OauthConnection | None": # A live match wins; among revoked matches, the latest consent wins. return ( - db_session() - .scalars( + await db_session().scalars( select(cls) .where( cls.provider == provider, @@ -132,29 +131,28 @@ def get_for_identity( .order_by(cls.revoked_at.is_(None).desc(), cls.connected_at.desc()) .limit(1) ) - .first() - ) + ).first() @classmethod - def list_for_provider( + async def list_for_provider( cls, provider: str, *, include_revoked: bool = False ) -> "list[OauthConnection]": query = select(cls).where(cls.provider == provider) if not include_revoked: query = query.where(cls.revoked_at.is_(None)) - return list(db_session().scalars(query)) + return list(await db_session().scalars(query)) @classmethod - def list_owned_by(cls, account_id: str | None) -> "list[OauthConnection]": + async def list_owned_by(cls, account_id: str | None) -> "list[OauthConnection]": # The audit read: everything this account ever authorized, revoked # rows included. return list( - db_session().scalars( + await db_session().scalars( select(cls).where(cls.account_id == account_id).order_by(cls.connected_at) ) ) - def reconnect( + async def reconnect( self, *, refresh_token: str, scopes: list[str], identity: dict[str, Any] | None = None ) -> None: self.refresh_token = refresh_token @@ -164,45 +162,43 @@ def reconnect( self.connected_at = Base.utc_now() self.revoked_at = None self.revoked_reason = "" - db_session().flush() + await db_session().flush() - def revoke(self, reason: str) -> None: + async def revoke(self, reason: str) -> None: # The consent's facts survive; only the secret is cleared. A second # revoke keeps the first stamp. self.revoked_at = self.revoked_at or Base.utc_now() self.revoked_reason = self.revoked_reason or reason self.refresh_token = "" - db_session().flush() + await db_session().flush() - def _load_refresh_token(self) -> str: + async def _load_refresh_token(self) -> str: # Under the refresh lock: another process may have rotated and # committed, and this transaction may already hold the row — # populate_existing re-reads it past the identity map. fresh = ( - db_session() - .scalars( + await db_session().scalars( select(OauthConnection) .where(OauthConnection.id == self.id) .execution_options(populate_existing=True) ) - .one() - ) + ).one() if fresh.revoked_at: raise OauthRefreshError(self.provider, "the connection was revoked mid-refresh") return fresh.refresh_token.decrypt() - def _save_refresh_token(self, rotated: str) -> None: + async def _save_refresh_token(self, rotated: str) -> None: # The provider invalidated the old token the moment it rotated, so # the write commits on its own session, never the enclosing # transaction — a step that rolls back later must not brick the # connection. - with get_session(db_session().get_bind()) as session: - stored = session.execute( + async with get_session(db_session().bind) as session: + stored = await session.execute( update(OauthConnection) .where(OauthConnection.id == self.id, OauthConnection.revoked_at.is_(None)) .values(refresh_token=rotated) ) - session.commit() + await session.commit() if not stored.rowcount: # A revoke landed mid-refresh. Nothing secret outlives the # consent at rest, so the rotated token is not stored. diff --git a/backend/druks/services/oauth.py b/backend/druks/services/oauth.py index 10575dea..5aa3d88e 100644 --- a/backend/druks/services/oauth.py +++ b/backend/druks/services/oauth.py @@ -217,7 +217,7 @@ async def get_access_token( try: data = { "grant_type": "refresh_token", - "refresh_token": connection._load_refresh_token(), + "refresh_token": await connection._load_refresh_token(), **self.extra_token_params, } if requested: @@ -250,7 +250,7 @@ async def get_access_token( self.provider, "the token endpoint returned no access token" ) if tokens.get("refresh_token"): - connection._save_refresh_token(tokens["refresh_token"]) + await connection._save_refresh_token(tokens["refresh_token"]) if requested and tokens.get("scope") and set(tokens["scope"].split()) != set(requested): # A provider that ignores the narrowing hands back a token the # sandbox must never hold — fail rather than cache it. @@ -280,7 +280,7 @@ async def evict_access_token(self, connection_id: str) -> None: async def disconnect(self, connection: OauthConnection, *, reason: str) -> None: """Revoke the connection and evict its cached access token. The row and its facts survive; the refresh token dies with the consent.""" - connection.revoke(reason) + await connection.revoke(reason) await self.evict_access_token(connection.id) diff --git a/backend/druks/services/routes.py b/backend/druks/services/routes.py index c560e0db..2a70a75c 100644 --- a/backend/druks/services/routes.py +++ b/backend/druks/services/routes.py @@ -25,13 +25,15 @@ async def list_services() -> list[ServiceResponse]: entries = [] for service in services.all(): try: - row = ServiceIdentity.get(service.slug) + row = await ServiceIdentity.get(service.slug) except ServiceNotConnectedError: row = None connections = [] if service.token_endpoint: # The detail shows revoked connections as history beside the live. - connections = OauthConnection.list_for_provider(service.slug, include_revoked=True) + connections = await OauthConnection.list_for_provider( + service.slug, include_revoked=True + ) entries.append(ServiceResponse.from_row(service, row, connections)) return entries @@ -56,7 +58,7 @@ async def connect_service(slug: str, payload: dict[str, str]) -> ServiceResponse # A replaced client can never refresh the old client's connections — # revoke every live one; the consents stay on record. client = OauthClient(provider=slug) - for connection in OauthConnection.list_for_provider(slug): + for connection in await OauthConnection.list_for_provider(slug): await client.disconnect(connection, reason="client_replaced") await publish( "oauth.disconnected", @@ -81,7 +83,7 @@ async def connect_oauth_service( service = _get_oauth_service(slug) account_id = current_account_id.get() if connection: - row = OauthConnection.get(connection) + row = await OauthConnection.get(connection) if not row or row.provider != slug: raise OauthPageError(f"No connection {connection!r} on {slug!r}.", status_code=404) if next and (not next.startswith("/") or next.startswith(("//", "/\\"))): @@ -95,7 +97,7 @@ async def connect_oauth_service( status_code=409, ) try: - client = service.get_oauth_client() + client = await service.get_oauth_client() except ServiceNotConnectedError as error: raise OauthPageError(str(error), status_code=409) from error url = await client.begin_connect( @@ -130,22 +132,24 @@ async def oauth_callback(state: str = "", code: str = "", error: str = "") -> Re # matches a fresh sign-in to one. Both make a revoked row live again. row = None if connection_id: - row = OauthConnection.get(connection_id) + row = await OauthConnection.get(connection_id) if not row: raise OauthPageError( "The connection was removed while consent was open.", status_code=400 ) elif service.identity_key and (value := identity.get(service.identity_key)): - row = OauthConnection.get_for_identity( + row = await OauthConnection.get_for_identity( provider, pending["account_id"], service.identity_key, value ) reconsent = bool(row) if row: - row.reconnect(refresh_token=tokens["refresh_token"], scopes=granted, identity=identity) + await row.reconnect( + refresh_token=tokens["refresh_token"], scopes=granted, identity=identity + ) # A token cached before this consent must not serve the new one. await OauthClient(provider=provider).evict_access_token(row.id) else: - row = OauthConnection.create( + row = await OauthConnection.create( provider=provider, account_id=pending["account_id"], refresh_token=tokens["refresh_token"], @@ -166,7 +170,7 @@ async def oauth_callback(state: str = "", code: str = "", error: str = "") -> Re @oauth_router.get("/connections", dependencies=[Depends(current_session_account)]) async def list_connections() -> list[ConnectionResponse]: - rows = OauthConnection.list_owned_by(current_account_id.get()) + rows = await OauthConnection.list_owned_by(current_account_id.get()) return [ConnectionResponse.model_validate(row) for row in rows] @@ -176,7 +180,7 @@ async def list_connections() -> list[ConnectionResponse]: dependencies=[Depends(current_session_account)], ) async def disconnect_connection(connection_id: str) -> None: - row = OauthConnection.get(connection_id) + row = await OauthConnection.get(connection_id) if not row: raise HTTPException(status_code=404, detail=f"No connection {connection_id!r}.") if row.revoked_at: diff --git a/backend/druks/signals.py b/backend/druks/signals.py index adfa0db4..767a2933 100644 --- a/backend/druks/signals.py +++ b/backend/druks/signals.py @@ -72,7 +72,7 @@ async def receiver(_sender: Any, **published: Any) -> None: return delivered = {key: value for key, value in published.items() if key not in _ROUTING} if subject_class: - subject = subject_class.get_for_subject_id(str(published["subject"]["id"])) + subject = await subject_class.get_for_subject_id(str(published["subject"]["id"])) if subject is None: return delivered["subject"] = subject diff --git a/backend/druks/skills/models.py b/backend/druks/skills/models.py index d396f709..9c57fca4 100644 --- a/backend/druks/skills/models.py +++ b/backend/druks/skills/models.py @@ -26,19 +26,22 @@ class SkillCollection(Base, Uuid7Pk): ) @classmethod - def list_all(cls) -> list["SkillCollection"]: - return list(db_session().execute(select(cls).order_by(cls.name)).scalars()) + async def list_all(cls) -> list["SkillCollection"]: + return list((await db_session().execute(select(cls).order_by(cls.name))).scalars()) @classmethod - def get(cls, collection_id: str) -> "SkillCollection | None": - return db_session().get(cls, collection_id) + async def get(cls, collection_id: str) -> "SkillCollection | None": + return await db_session().get(cls, collection_id) @classmethod - def get_for_source(cls, source: str) -> "SkillCollection | None": - return db_session().execute(select(cls).where(cls.source == source)).scalar_one_or_none() + async def get_for_source(cls, source: str) -> "SkillCollection | None": + result = await db_session().execute(select(cls).where(cls.source == source)) + return result.scalar_one_or_none() @classmethod - def create(cls, *, source: str, name: str, skills: list[InstalledSkill]) -> "SkillCollection": + async def create( + cls, *, source: str, name: str, skills: list[InstalledSkill] + ) -> "SkillCollection": session = db_session() collection = cls(source=source, name=name) collection.skills = [ @@ -51,13 +54,13 @@ def create(cls, *, source: str, name: str, skills: list[InstalledSkill]) -> "Ski for skill in skills ] session.add(collection) - session.flush() + await session.flush() return collection - def delete(self) -> None: + async def delete(self) -> None: session = db_session() - session.delete(self) - session.flush() + await session.delete(self) + await session.flush() class Skill(Base, Uuid7Pk): @@ -76,31 +79,32 @@ class Skill(Base, Uuid7Pk): updated_at: Mapped[datetime] = mapped_column(default=Base.utc_now, onupdate=Base.utc_now) @classmethod - def installed_names(cls) -> set[str]: - return set(db_session().execute(select(cls.name)).scalars()) + async def installed_names(cls) -> set[str]: + return set((await db_session().execute(select(cls.name))).scalars()) @classmethod - def list_enabled(cls) -> list["Skill"]: + async def list_enabled(cls) -> list["Skill"]: # The operator's enabled catalog. stmt = select(cls).where(cls.enabled.is_(True)).order_by(cls.name) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) @classmethod - def list_delivered(cls, requested: Collection[str]) -> list["Skill"]: + async def list_delivered(cls, requested: Collection[str]) -> list["Skill"]: # What one call receives: the enabled skills it named, or the whole enabled # catalog when it named none. - enabled = cls.list_enabled() + enabled = await cls.list_enabled() if requested: return [skill for skill in enabled if skill.name in requested] return enabled @classmethod - def delivery_excludes(cls, requested: Collection[str]) -> tuple[str, ...]: + async def delivery_excludes(cls, requested: Collection[str]) -> tuple[str, ...]: # Patterns are anchored to the skills_dir tar root (``-C skills_dir .``); # excluded skills remain installed on disk. - delivered = {skill.name for skill in cls.list_delivered(requested)} - return tuple(f"./{name}" for name in sorted(cls.installed_names() - delivered)) + delivered = {skill.name for skill in await cls.list_delivered(requested)} + return tuple(f"./{name}" for name in sorted(await cls.installed_names() - delivered)) @classmethod - def get(cls, name: str) -> "Skill | None": - return db_session().execute(select(cls).where(cls.name == name)).scalar_one_or_none() + async def get(cls, name: str) -> "Skill | None": + result = await db_session().execute(select(cls).where(cls.name == name)) + return result.scalar_one_or_none() diff --git a/backend/druks/skills/routes.py b/backend/druks/skills/routes.py index 862737fb..681823bd 100644 --- a/backend/druks/skills/routes.py +++ b/backend/druks/skills/routes.py @@ -13,7 +13,7 @@ @router.get("", response_model=list[CollectionResponse]) async def list_collections() -> list[SkillCollection]: - return SkillCollection.list_all() + return await SkillCollection.list_all() @router.post("", response_model=CollectionResponse) @@ -21,12 +21,12 @@ async def install_collection( settings: SettingsDep, url: str = Body(..., embed=True), ) -> SkillCollection: - if SkillCollection.get_for_source(url): + if await SkillCollection.get_for_source(url): raise HTTPException( status_code=409, detail=f"Collection {url!r} already installed; remove it first." ) try: - contents = await fetch_collection(url, settings.skills_dir, Skill.installed_names()) + contents = await fetch_collection(url, settings.skills_dir, await Skill.installed_names()) except (ValueError, RequestFailed, RequestTimeout) as error: raise HTTPException(status_code=422, detail=str(error)) from error except OSError as error: @@ -35,16 +35,16 @@ async def install_collection( raise HTTPException( status_code=500, detail=f"Could not write skills under {settings.skills_dir}: {error}" ) from error - return SkillCollection.create(source=url, name=contents.name, skills=contents.skills) + return await SkillCollection.create(source=url, name=contents.name, skills=contents.skills) @router.post("/{collection_id}/sync", response_model=CollectionResponse) async def sync_collection(collection_id: str, settings: SettingsDep) -> SkillCollection: - collection = SkillCollection.get(collection_id) + collection = await SkillCollection.get(collection_id) if not collection: raise HTTPException(status_code=404, detail=f"Collection {collection_id!r} not found") current_skills = {skill.name: skill for skill in collection.skills} - reserved_names = Skill.installed_names() - current_skills.keys() + reserved_names = await Skill.installed_names() - current_skills.keys() try: contents = await fetch_collection(collection.source, settings.skills_dir, reserved_names) except (ValueError, RequestFailed, RequestTimeout) as error: @@ -75,7 +75,7 @@ async def sync_collection(collection_id: str, settings: SettingsDep) -> SkillCol ) ) collection.updated_at = SkillCollection.utc_now() - db_session().flush() + await db_session().flush() return collection @@ -85,19 +85,19 @@ async def set_skill_enabled( name: str, enabled: bool = Body(..., embed=True), ) -> Skill: - skill = Skill.get(name) + skill = await Skill.get(name) if not skill or skill.collection_id != collection_id: raise HTTPException(status_code=404, detail=f"Skill {name!r} not found") skill.enabled = enabled - db_session().flush() + await db_session().flush() return skill @router.delete("/{collection_id}", status_code=204) async def remove_collection(collection_id: str) -> None: - collection = SkillCollection.get(collection_id) + collection = await SkillCollection.get(collection_id) if not collection: raise HTTPException(status_code=404, detail=f"Collection {collection_id!r} not found") for skill in collection.skills: remove_files(skill.path) - collection.delete() + await collection.delete() diff --git a/backend/druks/testing.py b/backend/druks/testing.py index 3dbc93d5..f978047d 100644 --- a/backend/druks/testing.py +++ b/backend/druks/testing.py @@ -3,16 +3,17 @@ import secrets import tempfile from collections.abc import AsyncIterator, Iterator +from contextlib import asynccontextmanager from pathlib import Path from unittest import mock +import httpx import pytest import redis.asyncio as aioredis from dbos import run_dbos_database_migrations -from fastapi.testclient import TestClient -from sqlalchemy import text, update +from sqlalchemy import NullPool, text, update from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from uuid_utils import uuid7 import druks.browser.models # noqa: F401 @@ -39,6 +40,7 @@ # fixtures the pytest11 entry point registers. Everything else here is the # repository suite's own harness. __all__ = [ + "asgi_client", "init_db", "run_workflow", "seed_call", @@ -54,6 +56,10 @@ ) TEST_REDIS_URL = os.environ.get("DRUKS_TEST_REDIS_URL", "redis://127.0.0.1:6379/15") +# The rollback fixture's async connection, for consumers that need the bind +# itself (the app under test binds its ambient session factory to it). +_fixture_connection = None + # Discovery runs once for the session, and a broken app is an ordinary state to # be in mid-edit. Holding the failure here lets a suite that never asks for a druks # fixture finish, and gives one that does an error naming the app. @@ -123,34 +129,40 @@ def _druks_schema(_druks_engine: Engine) -> Iterator[None]: @pytest.fixture -def druks_db(_druks_engine: Engine, _druks_schema: None) -> Iterator[Session]: - connection = _druks_engine.connect() - transaction = connection.begin() - # The app lifespan must not dispose the connection before this fixture rolls it back. - connection.dispose = lambda: None # type: ignore[method-assign] +async def druks_db(_druks_schema: None) -> AsyncIterator[AsyncSession]: + # Per-test async engine: connections bind to the running event loop, and + # pytest gives each test its own — NullPool so nothing pools across tests. + global _fixture_connection + engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool) + connection = await engine.connect() + transaction = await connection.begin() configure_session(connection) configure_engine(connection) - db_session.remove() - session = Session( + _fixture_connection = connection + session = AsyncSession( connection, join_transaction_mode="create_savepoint", autoflush=True, + expire_on_commit=False, ) db_session.registry.set(session) try: yield session finally: - db_session.remove() + _fixture_connection = None + await db_session.remove() configure_engine(None) - transaction.rollback() - connection.close() + if transaction.is_active: + await transaction.rollback() + await connection.close() + await engine.dispose() async def _operator_account(): from druks.accounts.context import current_account_id from druks.accounts.models import Account - account = Account.get_or_create("op@example.com") + account = await Account.get_or_create("op@example.com") current_account_id.set(account.id) return account @@ -180,7 +192,7 @@ def configure_app_for_test( from druks.api.server import app if not engine: - engine = db_session().get_bind() + engine = _fixture_connection configure_session(engine) app.state.settings = settings app.state.engine = engine @@ -190,21 +202,30 @@ def configure_app_for_test( return app +@asynccontextmanager +async def asgi_client(app) -> AsyncIterator[httpx.AsyncClient]: + """An HTTP client for ``app`` on the caller's own event loop — the test + fixtures' loop-bound connection rules out a portal-thread TestClient. No + lifespan: a fixture's teardown runs in another task, and the app lifespan's + anyio scopes cannot cross one; a test that needs the lifespan enters + ``app.router.lifespan_context(app)`` itself, inline.""" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + yield client + + @pytest.fixture -def druks_client(druks_db: Session, tmp_path: Path) -> Iterator[TestClient]: +async def druks_client(druks_db: AsyncSession, tmp_path: Path) -> AsyncIterator[httpx.AsyncClient]: from druks.accounts.dependencies import current_account, current_session_account dependencies = (current_account, current_session_account) - app = configure_app_for_test( - settings=make_settings(tmp_path), - engine=druks_db.connection(), - ) - # The app lifespan opens the Redis client on its own event loop and closes it - # there; one left behind by another loop cannot be closed from it. Drop it on - # both sides so each consumer dials on the loop it runs on. + app = configure_app_for_test(settings=make_settings(tmp_path)) + # ASGITransport, not TestClient: requests must run on this test's own event + # loop — the fixture connection is bound to it, and a portal thread's loop + # could never touch it. druks.redis._client = None try: - with TestClient(app) as client: + async with asgi_client(app) as client: yield client finally: druks.redis._client = None @@ -240,7 +261,7 @@ async def _phase_noop(*args, **kwargs): pass async def _cancel(workflow_id: str) -> None: - db_session().execute( + await db_session().execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == workflow_id) .values(status="CANCELLED") @@ -292,8 +313,8 @@ async def run_workflow( current_workflow.reset(token) -def seed_run( - session: Session, +async def seed_run( + session: AsyncSession, *, kind: str, subject: Subject | StoredSubject | None = None, @@ -315,17 +336,18 @@ def seed_run( account_id=account_id, ) session.add(run) - session.flush() + await session.flush() + await session.refresh(run, ["account"]) identity = None if subject: identity = {**subject.identity, "label": subject.label} - seed_dbos_status(session, run.id, state, subject=identity) - session.expire(run, ["state", "updated_at"]) + await seed_dbos_status(session, run.id, state, subject=identity) + await session.refresh(run, ["state", "updated_at"]) return run -def seed_dbos_status( - session: Session, +async def seed_dbos_status( + session: AsyncSession, run_id: str, state: str, *, @@ -350,7 +372,7 @@ def seed_dbos_status( # itself by its id — Subject.label's own rule. "subject_label": subject.get("label") or str(subject["id"]), } - session.execute( + await session.execute( workflow_status.insert().values( workflow_uuid=run_id, status=status, @@ -358,11 +380,11 @@ def seed_dbos_status( attributes=attributes, ) ) - session.flush() + await session.flush() -def seed_call( - session: Session, +async def seed_call( + session: AsyncSession, run: Run, agent: str, *, @@ -381,5 +403,8 @@ def seed_call( sandbox_host_id=f"test-host-{run.id}", ) session.add(call) - session.flush() + await session.flush() + # The read side serializes account and run off the row; a flush loads + # neither, and under async a lazy touch is an error. + await session.refresh(call, ["account", "run"]) return call diff --git a/backend/druks/usage/models.py b/backend/druks/usage/models.py index 56833a79..82def879 100644 --- a/backend/druks/usage/models.py +++ b/backend/druks/usage/models.py @@ -44,17 +44,19 @@ class UsageScrape(Base): unlimited: Mapped[bool] = mapped_column(default=False) @classmethod - def latest_for(cls, harness: str, account_id: str) -> "UsageScrape | None": + async def latest_for(cls, harness: str, account_id: str) -> "UsageScrape | None": stmt = ( select(cls) .where(cls.harness == harness, cls.account_id == account_id) .order_by(cls.scraped_at.desc()) .limit(1) ) - return db_session().execute(stmt).scalar_one_or_none() + return (await db_session().execute(stmt)).scalar_one_or_none() @classmethod - def history_for(cls, harness: str, account_id: str, *, since: datetime) -> list["UsageScrape"]: + async def history_for( + cls, harness: str, account_id: str, *, since: datetime + ) -> list["UsageScrape"]: """The account's successful scrapes for ``harness`` since ``since``, oldest first. Feeds the usage page's trend sparklines / burn-rate math, so failed scrapes (no percentages) are excluded.""" @@ -65,7 +67,7 @@ def history_for(cls, harness: str, account_id: str, *, since: datetime) -> list[ .where(cls.parse_ok.is_(True)) .order_by(cls.scraped_at.asc()) ) - return list(db_session().execute(stmt).scalars()) + return list((await db_session().execute(stmt)).scalars()) def binding_week(self) -> dict[str, Any] | None: """The window closest to exhaustion — whichever stops work first.""" @@ -85,18 +87,18 @@ def soonest_reset_after(self, now: datetime) -> datetime | None: if resets: return min(resets) - def save(self) -> None: + async def save(self) -> None: if not self.scraped_at: self.scraped_at = Base.utc_now() session = db_session() session.add(self) - session.flush() + await session.flush() @classmethod - def prune_older_than(cls, *, days: int) -> int: + async def prune_older_than(cls, *, days: int) -> int: cutoff = Base.utc_now() - timedelta(days=days) stmt = delete(cls).where(cls.scraped_at < cutoff) session = db_session() - result = session.execute(stmt) - session.flush() + result = await session.execute(stmt) + await session.flush() return result.rowcount diff --git a/backend/druks/usage/reads.py b/backend/druks/usage/reads.py index 9f56e5a0..25d825e0 100644 --- a/backend/druks/usage/reads.py +++ b/backend/druks/usage/reads.py @@ -6,20 +6,17 @@ from druks.durable.models import AgentCall -def list_finished_calls(account_id: str, *, since: datetime, until: datetime) -> list[Row]: - return list( - db_session() - .execute( - select( - AgentCall.model, - AgentCall.cost_usd, - AgentCall.cost_metadata, - AgentCall.finished_at, - ) - .where(AgentCall.account_id == account_id) - .where(AgentCall.finished_at.is_not(None)) - .where(AgentCall.finished_at >= since) - .where(AgentCall.finished_at < until) +async def list_finished_calls(account_id: str, *, since: datetime, until: datetime) -> list[Row]: + result = await db_session().execute( + select( + AgentCall.model, + AgentCall.cost_usd, + AgentCall.cost_metadata, + AgentCall.finished_at, ) - .all() + .where(AgentCall.account_id == account_id) + .where(AgentCall.finished_at.is_not(None)) + .where(AgentCall.finished_at >= since) + .where(AgentCall.finished_at < until) ) + return list(result.all()) diff --git a/backend/druks/usage/routes.py b/backend/druks/usage/routes.py index 7e46ecc5..7431041b 100644 --- a/backend/druks/usage/routes.py +++ b/backend/druks/usage/routes.py @@ -51,10 +51,10 @@ async def get_usage(account: Account = Depends(current_account)) -> UsageRespons now = datetime.now(UTC) summaries = [] for harness in get_harnesses(): - connection = HarnessConnection.get_for_account(harness.name, account.id) + connection = await HarnessConnection.get_for_account(harness.name, account.id) summaries.append( _summarize( - UsageScrape.latest_for(harness.name, account.id), + await UsageScrape.latest_for(harness.name, account.id), name=harness.name, now=now, connected=bool(connection), @@ -68,8 +68,8 @@ async def get_usage(account: Account = Depends(current_account)) -> UsageRespons async def refresh_usage(account: Account = Depends(current_account)) -> None: now = datetime.now(UTC) for harness in get_harnesses(): - connection = HarnessConnection.get_for_account(harness.name, account.id) - row = UsageScrape.latest_for(harness.name, account.id) + connection = await HarnessConnection.get_for_account(harness.name, account.id) + row = await UsageScrape.latest_for(harness.name, account.id) age = _age_seconds(row.scraped_at, now=now) if row else None if connection and (age is None or age >= _REFRESH_FLOOR_SECONDS): await harness.poll_usage(connection) @@ -83,7 +83,7 @@ async def refresh_usage(account: Account = Depends(current_account)) -> None: async def get_usage_history(account: Account = Depends(current_account)) -> UsageHistoryResponse: now = datetime.now(UTC) return UsageHistoryResponse( - harnesses=[_harness_history(h.name, account.id, now=now) for h in get_harnesses()], + harnesses=[await _harness_history(h.name, account.id, now=now) for h in get_harnesses()], ) @@ -95,8 +95,12 @@ async def get_usage_history(account: Account = Depends(current_account)) -> Usag async def get_usage_today(account: Account = Depends(current_account)) -> UsageTodayResponse: # Deriving the operator-local-day window here (the query just takes it) keeps # this total identical to the sys-strip's and the agent surface's figures. - timezone, local_start = operator_local_day(UserSettings.get().timezone, datetime.now(UTC)) - rows = list_finished_calls(account.id, since=local_start, until=local_start + timedelta(days=1)) + timezone, local_start = operator_local_day( + (await UserSettings.get()).timezone, datetime.now(UTC) + ) + rows = await list_finished_calls( + account.id, since=local_start, until=local_start + timedelta(days=1) + ) timezone_name = str(timezone) # Every call counts, even one whose model no picker list claims (pinned @@ -107,7 +111,7 @@ async def get_usage_today(account: Account = Depends(current_account)) -> UsageT # whole list. harness_by_model = { entry["id"]: settings.name - for settings in HarnessSettings.all() + for settings in await HarnessSettings.all() for entry in settings.allowed_models } names = [h.name for h in get_harnesses()] @@ -141,8 +145,8 @@ async def get_usage_today(account: Account = Depends(current_account)) -> UsageT ) -def _harness_history(name: str, account_id: str, *, now: datetime) -> UsageHarnessHistory: - rows = UsageScrape.history_for(name, account_id, since=now - WEEK_RANGE) +async def _harness_history(name: str, account_id: str, *, now: datetime) -> UsageHarnessHistory: + rows = await UsageScrape.history_for(name, account_id, since=now - WEEK_RANGE) five_hour_cutoff = now - FIVE_HOUR_RANGE five_hour = [ UsageHistoryPoint(t=row.scraped_at, pct=row.five_hour_percent_left) diff --git a/backend/druks/user_settings/models.py b/backend/druks/user_settings/models.py index 84d0133b..1fa80a44 100644 --- a/backend/druks/user_settings/models.py +++ b/backend/druks/user_settings/models.py @@ -45,30 +45,32 @@ class UserSettings(Base): SINGLETON_ID = 1 @classmethod - def get(cls) -> "UserSettings": + async def get(cls) -> "UserSettings": session = db_session() - row = session.get(cls, cls.SINGLETON_ID) + row = await session.get(cls, cls.SINGLETON_ID) if not row: - session.execute(pg_insert(cls).values(id=cls.SINGLETON_ID).on_conflict_do_nothing()) - row = session.get_one(cls, cls.SINGLETON_ID) + await session.execute( + pg_insert(cls).values(id=cls.SINGLETON_ID).on_conflict_do_nothing() + ) + row = await session.get_one(cls, cls.SINGLETON_ID) return row - def update_profile(self, *, timezone: str | None = None) -> None: + async def update_profile(self, *, timezone: str | None = None) -> None: if timezone: self.timezone = timezone self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def set_fallback_account(self, account_id: str) -> None: + async def set_fallback_account(self, account_id: str) -> None: self.fallback_account_id = account_id self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() - def set_gate_park_destination(self, destination_id: str | None) -> None: + async def set_gate_park_destination(self, destination_id: str | None) -> None: # None is the off-switch, so this is a set-or-clear, not a skip-on-None. self.gate_park_destination_id = destination_id self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() class HarnessSettings(Base): @@ -90,21 +92,21 @@ class HarnessSettings(Base): models_fetched_at: Mapped[datetime | None] = mapped_column(default=None) @classmethod - def get(cls, name: str) -> "HarnessSettings | None": - return db_session().get(cls, name) + async def get(cls, name: str) -> "HarnessSettings | None": + return await db_session().get(cls, name) @classmethod - def require(cls, name: str) -> "HarnessSettings": + async def require(cls, name: str) -> "HarnessSettings": # The resolution paths (effort/timeout, the harness factory) only pass a # ``get_harness_for_model`` name, which is always registered and so always # seeded — a miss means ``seed_harnesses`` didn't run before serving. - if not (config := cls.get(name)): + if not (config := await cls.get(name)): raise KeyError(f"no harness settings for {name!r}; seed_harnesses missed it") return config @classmethod - def all(cls) -> list["HarnessSettings"]: - return list(db_session().execute(select(cls).order_by(cls.name)).scalars()) + async def all(cls) -> list["HarnessSettings"]: + return list((await db_session().execute(select(cls).order_by(cls.name))).scalars()) @property def harness(self) -> "type[Harness]": @@ -136,16 +138,16 @@ async def refresh_models(self, connection: "HarnessConnection") -> dict[str, obj if parsed.ok: self.models_fetched = list(parsed.models) self.models_fetched_at = Base.utc_now() - db_session().flush() + await db_session().flush() return {"harness": self.name, "ok": parsed.ok, "error": parsed.error} - def update(self, **fields: object) -> None: + async def update(self, **fields: object) -> None: # Callers pass column names only — the route's ``HarnessUpdate`` schema is # the trust boundary, so no field-name validation here. for field, value in fields.items(): setattr(self, field, value) self.updated_at = Base.utc_now() - db_session().flush() + await db_session().flush() class SettingsOverride(Base): @@ -156,94 +158,95 @@ class SettingsOverride(Base): secret_value = EncryptedTextField(default="") @classmethod - def read(cls, key: str) -> Any | None: - row = db_session().get(cls, key) + async def read(cls, key: str) -> Any | None: + row = await db_session().get(cls, key) return row.value if row else None @classmethod - def write(cls, key: str, value: Any) -> None: + async def write(cls, key: str, value: Any) -> None: session = db_session() - row = session.get(cls, key) + row = await session.get(cls, key) if value is None: if row: - session.delete(row) + await session.delete(row) elif row: row.value = value else: session.add(cls(key=key, value=value)) - session.flush() + await session.flush() @classmethod - def agent_model(cls, name: str, default: str) -> ResolvedModel: - override = cls.read(f"agent_model:{name}") + async def agent_model(cls, name: str, default: str) -> ResolvedModel: + override = await cls.read(f"agent_model:{name}") if override is not None: return ResolvedModel(override, "agent") # ``default`` is a harness name (claude/codex) → that harness's model, # or a pinned model string when it names no harness. - harness = HarnessSettings.get(default) + harness = await HarnessSettings.get(default) return ResolvedModel(harness.model if harness else default, "default") @classmethod - def set_agent_model(cls, name: str, model: str | None) -> None: - cls.write(f"agent_model:{name}", model) + async def set_agent_model(cls, name: str, model: str | None) -> None: + await cls.write(f"agent_model:{name}", model) @classmethod - def agent_effort(cls, name: str, declared: str | None, harness: str) -> ResolvedEffort: - override = cls.read(f"agent_effort:{name}") + async def agent_effort(cls, name: str, declared: str | None, harness: str) -> ResolvedEffort: + override = await cls.read(f"agent_effort:{name}") if override is not None: return ResolvedEffort(override, "agent") if declared is not None: return ResolvedEffort(declared, "declared") - return ResolvedEffort(HarnessSettings.require(harness).effort, "harness") + return ResolvedEffort((await HarnessSettings.require(harness)).effort, "harness") @classmethod - def set_agent_effort(cls, name: str, value: str | None) -> None: - cls.write(f"agent_effort:{name}", value) + async def set_agent_effort(cls, name: str, value: str | None) -> None: + await cls.write(f"agent_effort:{name}", value) @classmethod - def agent_timeout(cls, name: str, declared: int | None, harness: str) -> ResolvedTimeout: - override = cls.read(f"agent_timeout:{name}") + async def agent_timeout(cls, name: str, declared: int | None, harness: str) -> ResolvedTimeout: + override = await cls.read(f"agent_timeout:{name}") if override is not None: return ResolvedTimeout(override, "agent") if declared is not None: return ResolvedTimeout(declared, "declared") - return ResolvedTimeout(HarnessSettings.require(harness).timeout, "harness") + return ResolvedTimeout((await HarnessSettings.require(harness)).timeout, "harness") @classmethod - def set_agent_timeout(cls, name: str, value: int | None) -> None: - cls.write(f"agent_timeout:{name}", value) + async def set_agent_timeout(cls, name: str, value: int | None) -> None: + await cls.write(f"agent_timeout:{name}", value) @classmethod - def workflow_setting(cls, kind: str, field: str, default: Any) -> Any: - value = cls.read(f"workflow:{kind}:{field}") + async def workflow_setting(cls, kind: str, field: str, default: Any) -> Any: + value = await cls.read(f"workflow:{kind}:{field}") return default if value is None else value @classmethod - def set_workflow_setting(cls, kind: str, field: str, value: Any) -> None: - cls.write(f"workflow:{kind}:{field}", value) + async def set_workflow_setting(cls, kind: str, field: str, value: Any) -> None: + await cls.write(f"workflow:{kind}:{field}", value) @classmethod - def app_setting(cls, app: str, field: str, default: Any, *, is_secret: bool) -> Any: - row = db_session().get(cls, f"app:{app}:{field}") + async def app_setting(cls, app: str, field: str, default: Any, *, is_secret: bool) -> Any: + row = await db_session().get(cls, f"app:{app}:{field}") if is_secret: return row.secret_value.decrypt() if row and row.secret_value else default return row.value if row else default @classmethod - def set_app_setting(cls, app: str, field: str, value: Any, *, is_secret: bool) -> None: + async def set_app_setting(cls, app: str, field: str, value: Any, *, is_secret: bool) -> None: key = f"app:{app}:{field}" if value is None or not is_secret: - cls.write(key, value) + await cls.write(key, value) return session = db_session() - row = session.get(cls, key) + row = await session.get(cls, key) if row: row.value = None row.secret_value = value else: row = cls(key=key, value=None, secret_value=value) session.add(row) - session.flush() - # Assignment leaves the plaintext str on the instance; expire it so the - # next read loads the envelope. - session.expire(row) + await session.flush() + # Assignment leaves the plaintext str on the instance; reload it now so + # the next read sees the envelope — an expired attribute can't + # lazy-load under the async session. + await session.refresh(row) diff --git a/backend/druks/user_settings/reads.py b/backend/druks/user_settings/reads.py index 9d041ff8..6e4feb66 100644 --- a/backend/druks/user_settings/reads.py +++ b/backend/druks/user_settings/reads.py @@ -22,11 +22,11 @@ from druks.workflows import Workflow -def get_agent_setting(agent: "Agent") -> AgentSettingResponse: - model = SettingsOverride.agent_model(agent.id, agent.model) - harness = get_harness_for_model(model.value).name - effort = SettingsOverride.agent_effort(agent.id, agent.effort, harness) - timeout = SettingsOverride.agent_timeout(agent.id, agent.timeout, harness) +async def get_agent_setting(agent: "Agent") -> AgentSettingResponse: + model = await SettingsOverride.agent_model(agent.id, agent.model) + harness = (await get_harness_for_model(model.value)).name + effort = await SettingsOverride.agent_effort(agent.id, agent.effort, harness) + timeout = await SettingsOverride.agent_timeout(agent.id, agent.timeout, harness) return AgentSettingResponse( name=agent.name or agent.id, description=agent.description, @@ -40,20 +40,20 @@ def get_agent_setting(agent: "Agent") -> AgentSettingResponse: ) -def get_settings_field( +async def get_settings_field( name: str, field: FieldInfo, *, value: Any, override_key: str ) -> SettingsFieldResponse: - overridden = db_session().get(SettingsOverride, override_key) is not None + overridden = await db_session().get(SettingsOverride, override_key) is not None return SettingsFieldResponse.from_field(name, field, value=value, overridden=overridden) -def get_workflow_settings(workflow: "type[Workflow]") -> WorkflowSettingsResponse: +async def get_workflow_settings(workflow: "type[Workflow]") -> WorkflowSettingsResponse: kind = workflow.kind fields = [ - get_settings_field( + await get_settings_field( name, field, - value=SettingsOverride.workflow_setting(kind, name, field.default), + value=await SettingsOverride.workflow_setting(kind, name, field.default), override_key=f"workflow:{kind}:{name}", ) for name, field in workflow.Settings.model_fields.items() @@ -70,53 +70,54 @@ def get_workflow_settings(workflow: "type[Workflow]") -> WorkflowSettingsRespons # "cron" is a UI kind like enum/secret: the frontend renders # cadence presets with a raw-cron escape hatch. type="cron", - value=workflow.get_schedule(), + value=await workflow.get_schedule(), default=workflow.every, choices=None, section="", visible_when_field="", visible_when_value=None, secret_set=None, - overridden=SettingsOverride.read(f"workflow:{kind}:schedule") is not None, + overridden=await SettingsOverride.read(f"workflow:{kind}:schedule") is not None, ), SettingsFieldResponse( name="schedule_enabled", label=f"{label} enabled", help="Pause the scheduled run without losing its cadence.", type="bool", - value=workflow.has_enabled_schedule(), + value=await workflow.has_enabled_schedule(), default=True, choices=None, section="", visible_when_field="", visible_when_value=None, secret_set=None, - overridden=SettingsOverride.read(f"workflow:{kind}:schedule_enabled") is not None, + overridden=await SettingsOverride.read(f"workflow:{kind}:schedule_enabled") + is not None, ), ] return WorkflowSettingsResponse(kind=kind, fields=fields) -def get_app_settings(app: "type[App]") -> AppSettingsResponse: +async def get_app_settings(app: "type[App]") -> AppSettingsResponse: model = app.settings_model return AppSettingsResponse( name=app.name, description=app.description, icon=app.icon, builtin=app.builtin, - agents=[get_agent_setting(agent) for agent in app.agents()], + agents=[await get_agent_setting(agent) for agent in app.agents()], # Surface only the workflows with operator knobs: tunable settings or a # schedule to retune. workflows=[ - get_workflow_settings(workflow) + await get_workflow_settings(workflow) for workflow in app.workflows() if workflow.Settings.model_fields or workflow.every ], settings=[ - get_settings_field( + await get_settings_field( name, field, - value=SettingsOverride.app_setting( + value=await SettingsOverride.app_setting( app.name, name, field.default, diff --git a/backend/druks/user_settings/routes.py b/backend/druks/user_settings/routes.py index 7442915b..a4155bed 100644 --- a/backend/druks/user_settings/routes.py +++ b/backend/druks/user_settings/routes.py @@ -38,15 +38,15 @@ def _validate_timezone(value: str) -> str: return value -def _resolve_harness(name: str) -> tuple[type, HarnessSettings]: - row = HarnessSettings.get(name) +async def _resolve_harness(name: str) -> tuple[type, HarnessSettings]: + row = await HarnessSettings.get(name) if row: return row.harness, row raise HTTPException(status_code=404, detail=f"Unknown harness: {name!r}") -def _harness_response(settings: HarnessSettings, account: Account) -> HarnessResponse: - connection = HarnessConnection.get_for_account(settings.name, account.id) +async def _harness_response(settings: HarnessSettings, account: Account) -> HarnessResponse: + connection = await HarnessConnection.get_for_account(settings.name, account.id) return HarnessResponse.from_row(settings, connection, account) @@ -58,7 +58,9 @@ async def list_harness_settings( ) -> list[HarnessResponse]: registered = {harness.name for harness in get_harnesses()} return [ - _harness_response(row, account) for row in HarnessSettings.all() if row.name in registered + await _harness_response(row, account) + for row in await HarnessSettings.all() + if row.name in registered ] @@ -66,11 +68,11 @@ async def list_harness_settings( async def update_harness_settings( name: str, body: HarnessUpdate, account: Account = Depends(current_account) ) -> HarnessResponse: - harness, row = _resolve_harness(name) + harness, row = await _resolve_harness(name) updates = body.model_dump(exclude_unset=True, by_alias=False) if "model" in updates: try: - resolved = get_harness_for_model(updates["model"]) + resolved = await get_harness_for_model(updates["model"]) if resolved.name != harness.name: raise HarnessError except HarnessError as exc: @@ -81,38 +83,38 @@ async def update_harness_settings( _validate_effort(updates.get("effort")) _validate_timeout(updates.get("timeout")) if updates: - row.update(**updates) - return _harness_response(row, account) + await row.update(**updates) + return await _harness_response(row, account) @router.get("", response_model=UserSettingsResponse, response_model_by_alias=True) async def get_user_settings() -> UserSettings: - return UserSettings.get() + return await UserSettings.get() @router.patch("", response_model=UserSettingsResponse, response_model_by_alias=True) async def update_user_settings( body: UpdateUserSettingsRequest, ) -> UserSettings: - row = UserSettings.get() + row = await UserSettings.get() if body.timezone is not None: tz = _validate_timezone(body.timezone) if tz != row.timezone: - row.update_profile(timezone=tz) + await row.update_profile(timezone=tz) # Crons are evaluated in this timezone — repoint them now, not at # the next launch. - apply_schedules() + await apply_schedules() if "gate_park_destination_id" in body.model_fields_set: destination_id = body.gate_park_destination_id - if destination_id and not Destination.get(destination_id): + if destination_id and not await Destination.get(destination_id): raise HTTPException(status_code=422, detail=f"Unknown destination {destination_id!r}") - row.set_gate_park_destination(destination_id) + await row.set_gate_park_destination(destination_id) return row @router.get("/apps", response_model=AppsSettingsResponse, response_model_by_alias=True) async def get_app_settings() -> AppsSettingsResponse: - projected = (reads.get_app_settings(m) for m in iter_apps()) + projected = [await reads.get_app_settings(m) for m in iter_apps()] return AppsSettingsResponse( allowed_efforts=list(ALLOWED_EFFORTS), apps=[out for out in projected if out.agents or out.workflows or out.settings], @@ -121,11 +123,11 @@ async def get_app_settings() -> AppsSettingsResponse: # An agent's model override is client data — reject a model no installed harness # lists. -def _validate_model(value: str | None) -> None: +async def _validate_model(value: str | None) -> None: if value is None: return try: - get_harness_for_model(value) + await get_harness_for_model(value) except HarnessError as exc: raise HTTPException( status_code=422, @@ -157,16 +159,16 @@ def _validate_timeout(value: int | None) -> None: ) async def update_app_settings(body: AppsSettingsUpdate) -> AppsSettingsResponse: for name, model in body.agent_models.items(): - _validate_model(model) - SettingsOverride.set_agent_model(name, model) + await _validate_model(model) + await SettingsOverride.set_agent_model(name, model) for name, effort in body.agent_efforts.items(): _validate_effort(effort) - SettingsOverride.set_agent_effort(name, effort) + await SettingsOverride.set_agent_effort(name, effort) for name, timeout in body.agent_timeouts.items(): _validate_timeout(timeout) - SettingsOverride.set_agent_timeout(name, timeout) + await SettingsOverride.set_agent_timeout(name, timeout) changed_apps = [] try: @@ -175,14 +177,14 @@ async def update_app_settings(body: AppsSettingsUpdate) -> AppsSettingsResponse: if not workflow: raise HTTPException(status_code=422, detail=f"Unknown workflow {kind!r}") for field, value in changes.items(): - workflow.override_setting(field, value) + await workflow.override_setting(field, value) for app_name, changes in body.app_settings.items(): try: app = get_app(app_name) except KeyError as exc: raise HTTPException(status_code=422, detail=f"Unknown app {app_name!r}") from exc for field, value in changes.items(): - app.override_setting(field, value) + await app.override_setting(field, value) changed_apps.append(app) except ValueError as exc: # Domain rejections (unknown field, bad cron, failed constraint) → 422. @@ -192,7 +194,7 @@ async def update_app_settings(body: AppsSettingsUpdate) -> AppsSettingsResponse: settings_problems = {} for app in changed_apps: - if problems := app.settings().clean(): + if problems := (await app.settings()).clean(): settings_problems[app.name] = problems if settings_problems: raise HTTPException(status_code=422, detail=settings_problems) @@ -204,6 +206,6 @@ async def update_app_settings(body: AppsSettingsUpdate) -> AppsSettingsResponse: ): # Repoint the DBOS crons now, not at the next launch; the reconcile reads # the just-written overrides off this request's session. - apply_schedules() + await apply_schedules() return await get_app_settings() diff --git a/backend/druks/webhooks/base.py b/backend/druks/webhooks/base.py index 34f4dd28..5f9448a7 100644 --- a/backend/druks/webhooks/base.py +++ b/backend/druks/webhooks/base.py @@ -95,7 +95,7 @@ def get_action(self) -> str: """ raise NotImplementedError - def request_is_authentic(self) -> bool: + async def request_is_authentic(self) -> bool: """Verify the request is from the claimed provider. Return ``True`` if the request should be processed, ``False`` to @@ -145,7 +145,7 @@ def log_ignored(self, *, event: str, reason: str, **extra: Any) -> None: async def respond(self) -> Response: self.raw_body = await self.request.body() - if not self.request_is_authentic(): + if not await self.request_is_authentic(): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid webhook signature.", diff --git a/backend/druks/workflows.py b/backend/druks/workflows.py index 4a0a2a18..8c17904b 100644 --- a/backend/druks/workflows.py +++ b/backend/druks/workflows.py @@ -156,11 +156,23 @@ def __init__(self, subject_class: "type[Subject] | type[StoredSubject] | None") def __get__(self, run: "Workflow | None", owner: type) -> Any: if run is None: return self.subject_class + # Live, not a snapshot taken at dispatch: a long-parked run resumes against # whatever the declared class says then, and finds nothing if it went away. - if not run._subject: - return - return self.subject_class.get_for_subject_id(str(run._subject["id"])) + # Awaitable either way, so ``await self.subject`` is the one shape. + async def resolve() -> Any: + if "subject" in run.__dict__: + return run.__dict__["subject"] + if not run._subject: + return None + return await self.subject_class.get_for_subject_id(str(run._subject["id"])) + + return resolve() + + def __set__(self, run: "Workflow", value: Any) -> None: + # A test hands the run its subject directly; ``await self.subject`` + # then resolves to it without a read. + run.__dict__["subject"] = value def _declare_subject(cls: type["Workflow"]) -> None: @@ -267,7 +279,7 @@ async def answer(cls, subject: Subject | StoredSubject, **reply: Any) -> None: f"{cls.__name__}.answer() takes the subject whose run is parked on it, " f"not {type(subject).__name__}" ) - runs = Run.list_for_subject(subject.subject_type, str(subject.id)) + runs = await Run.list_for_subject(subject.subject_type, str(subject.id)) parked = next((run for run in runs if run.is_parked and run.input_gate == cls.name), None) if parked: await parked.resume(**reply) @@ -356,9 +368,10 @@ async def _notify_designated_destination(workflow_id: str, subject: dict[str, An # the settings pointer is the operator's off-switch. async def _create() -> str | None: async with step_session(): - destination_id = UserSettings.get().gate_park_destination_id + destination_id = (await UserSettings.get()).gate_park_destination_id if destination_id: - return Run.get(workflow_id).create_park_notification(destination_id, subject) + run = await Run.get(workflow_id) + return await run.create_park_notification(destination_id, subject) notification_id = await DBOS.run_step_async( StepOptions(name="notifications.gate_park", **_IO_RETRIES), _create @@ -532,17 +545,20 @@ async def _emit_run_event( # own arguments, so a replay stamps the same routing every time. async def _transition() -> dict[str, Any] | None: async with step_session() as session: - run = Run.get(workflow_id) + run = await Run.get(workflow_id) + # Read before the flush: flushing the update unloads the row's + # computed columns, and reading one back would be implicit IO. + label = run.subject_label if facts: for field, value in facts.items(): setattr(run, field, value) - session.flush() + await session.flush() # Subjectless framework crons are plumbing: no feed entry. if subject: return { "kind": run.kind, "subject": subject, - "payload": _log_run_event(run, state, subject, result), + "payload": await _log_run_event(run, state, subject, label, result), } transition = await DBOS.run_step_async( @@ -565,10 +581,11 @@ async def _propagate() -> None: ) -def _log_run_event( +async def _log_run_event( run: Run, state: RunState, subject: dict[str, Any], + label: str | None, result: Any = None, ) -> dict[str, Any]: # One event per transition — the feed's run-level granularity, read off the @@ -584,10 +601,10 @@ def _log_run_event( payload["result"] = result.model_dump(mode="json") elif isinstance(result, dict): payload["result"] = result - Event.emit( + await Event.emit( type=WorkflowEvent.for_state(state), subject=subject, - label=run.subject_label, + label=label, payload=payload, app=workflows.get(run.kind).app, ) @@ -625,7 +642,7 @@ async def _execute_run( # Every failure re-raises so DBOS records the terminal ERROR derived state # reads; an operator cancel already carries its own reason and terminal # status, so it passes through untouched. - Run.create_row(_step_engine(), workflow_id=workflow_id, kind=kind, account_id=account_id) + await Run.create_row(_step_engine(), workflow_id=workflow_id, kind=kind, account_id=account_id) await _emit_run_event(workflow_id, RunState.RUNNING, subject=subject) async def record_failed(exc: BaseException, code: str) -> None: @@ -844,28 +861,28 @@ def workflow_id(self) -> str: # operator override → the declared default. The reconciler and the settings # read both go through these, so the workflow owns its own knobs. @classmethod - def get_schedule(cls) -> str | None: - return SettingsOverride.workflow_setting(cls.kind, "schedule", cls.every) + async def get_schedule(cls) -> str | None: + return await SettingsOverride.workflow_setting(cls.kind, "schedule", cls.every) @classmethod - def has_enabled_schedule(cls) -> bool: + async def has_enabled_schedule(cls) -> bool: # There is a schedule and it's on — False for unscheduled workflows too. if not cls.every: return False - return SettingsOverride.workflow_setting(cls.kind, "schedule_enabled", True) + return await SettingsOverride.workflow_setting(cls.kind, "schedule_enabled", True) @classmethod - def settings(cls) -> BaseModel: + async def settings(cls) -> BaseModel: """The workflow's ``Settings``, resolved through the override store — the read twin of ``override_setting``, like ``App.settings()`` for an app.""" values = { - name: SettingsOverride.workflow_setting(cls.kind, name, field.default) + name: await SettingsOverride.workflow_setting(cls.kind, name, field.default) for name, field in cls.Settings.model_fields.items() } return cls.Settings.model_validate(values) @classmethod - def override_setting(cls, field: str, value: Any) -> None: + async def override_setting(cls, field: str, value: Any) -> None: # An operator's override for one knob; None clears it back to the declared # default. Raises ValueError so the API layer can 422 it. The schedule pair # is validated here, not against Settings — those knobs live beside every=. @@ -878,8 +895,10 @@ def override_setting(cls, field: str, value: Any) -> None: raise ValueError(f"Unknown {cls.kind} setting {field!r}") elif value is not None: value = coerce_setting_value(cls.Settings, field, value) - validate_setting_override(cls.Settings, cls.settings().model_dump(), field, value) - SettingsOverride.set_workflow_setting(cls.kind, field, value) + validate_setting_override( + cls.Settings, (await cls.settings()).model_dump(), field, value + ) + await SettingsOverride.set_workflow_setting(cls.kind, field, value) @classmethod def _validate_subject(cls, subject: "Subject | StoredSubject | None") -> None: @@ -900,7 +919,7 @@ def _validate_subject(cls, subject: "Subject | StoredSubject | None") -> None: @classmethod async def cancel(cls, subject: Subject | StoredSubject, *, failure: str | None = None) -> None: cls._validate_subject(subject) - runs = Run.list_for_subject(subject.subject_type, str(subject.id), kind=cls.kind) + runs = await Run.list_for_subject(subject.subject_type, str(subject.id), kind=cls.kind) run = next((run for run in runs if run.is_active), None) if run: await run.cancel(failure=failure) @@ -970,7 +989,7 @@ async def start( if handle.workflow_id == workflow_id: # The body also creates its row (idempotently) — this one just makes it # visible before an executor picks the workflow up. - Run.create_row( + await Run.create_row( _step_engine(), workflow_id=workflow_id, kind=cls.kind, account_id=account_id ) if subject: diff --git a/backend/druks/workspaces.py b/backend/druks/workspaces.py index e333c2ce..ad519737 100644 --- a/backend/druks/workspaces.py +++ b/backend/druks/workspaces.py @@ -46,7 +46,7 @@ async def run_agent(self, *, account_id: str | None, **kwargs: Any) -> AgentResu run_kwargs = await self.with_mcp_servers(account_id, **self.get_agent_run_kwargs(**kwargs)) # with_mcp_servers is the run's last DB read; commit so the step's # connection isn't held idle through the minutes the agent runs. - db_session().commit() + await db_session().commit() return await self.sandbox.run_agent(**run_kwargs) async def with_mcp_servers(self, account_id: str | None, **kwargs: Any) -> dict[str, Any]: @@ -60,10 +60,10 @@ async def with_mcp_servers(self, account_id: str | None, **kwargs: Any) -> dict[ # One config key per name in the emitted harness config — a dupe # would break the VM's config parse mid-run. raise ValueError(f"duplicate required MCP server names: {sorted(required_names)}") - enabled = mcp_models.McpServer.list_enabled() + enabled = await mcp_models.McpServer.list_enabled() if not required and not enabled: return kwargs - run_account = account_id or UserSettings.get().fallback_account_id + run_account = account_id or (await UserSettings.get()).fallback_account_id # ``extra_env`` may be omitted or an explicit ``None`` (both valid for the # underlying run_agent); treat them the same so the merge never unpacks None. env = dict(kwargs.get("extra_env") or {}) @@ -152,14 +152,14 @@ async def set_git_identity(self, account_id: str | None) -> None: Rewritten before every agent call so a reused warm host follows the current run's dispatcher — a system dispatch carries no hook and credits nobody.""" - author_name, author_email = await get_github_client().get_bot_git_author() + author_name, author_email = await (await get_github_client()).get_bot_git_author() steps = [ f"cd {shlex.quote(self.repo_path)}", f"git config user.name {shlex.quote(author_name)}", f"git config user.email {shlex.quote(author_email)}", "rm -f .git/hooks/prepare-commit-msg", ] - if account_id and (account := Account.get(account_id, exclude_system=True)): + if account_id and (account := await Account.get(account_id, exclude_system=True)): trailer = f"Co-Authored-By: {account.username} <{account.username}>" hook = ( "#!/bin/sh\n" diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 669c28cb..51c1491f 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -48,7 +48,7 @@ async def _dbos_cancel(workflow_id: str) -> None: from druks.durable.dbos_state import workflow_status from sqlalchemy import update - db_session().execute( + await db_session().execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == workflow_id) .values(status="CANCELLED") @@ -110,24 +110,26 @@ def browser_session_declarations(): # opt out and reset themselves. Everything else — including the durable *unit* # tests that use the fixtures here — gets transaction rollback. _OWN_DATABASE_MODULES = { - "test_build_durable", "test_durable_sdk", "test_notifications_durable", "test_harness_login_persistence", "test_app_migrations", - "test_plan_gate_migration", "test_proof_app_migration", } -@pytest.fixture(autouse=True) -def _platform_database(request): - if request.module.__name__.rsplit(".", 1)[-1] in _OWN_DATABASE_MODULES: - yield - return - - request.getfixturevalue("druks_db") - yield +def pytest_collection_modifyitems(items): + # druks_db is async, so a sync fixture can't getfixturevalue it any more; + # injecting it into each test's fixture list keeps the same guarantee — + # every test outside the own-database modules runs inside the rollback + # transaction. + for item in items: + module = item.module.__name__.rsplit(".", 1)[-1] + if module in _OWN_DATABASE_MODULES: + continue + if not hasattr(item, "fixturenames") or "druks_db" in item.fixturenames: + continue + item.fixturenames.append("druks_db") @pytest.fixture(autouse=True) @@ -160,18 +162,18 @@ def bind_ambient_session(session) -> None: db_session.registry.set(session) -def connect_harness(harness_cls, payload: dict, *, provider_email: str = "op@example.com"): +async def connect_harness(harness_cls, payload: dict, *, provider_email: str = "op@example.com"): """Seed the HarnessConnection row a finished connect flow would leave.""" from druks.accounts.models import Account from druks.harnesses.models import HarnessConnection from druks.user_settings.models import UserSettings - account = Account.get_or_create(provider_email) - settings = UserSettings.get() + account = await Account.get_or_create(provider_email) + settings = await UserSettings.get() if not settings.fallback_account_id: - settings.set_fallback_account(account.id) + await settings.set_fallback_account(account.id) _, expires_at = harness_cls._refresh_state(payload) - return HarnessConnection.connect( + return await HarnessConnection.connect( harness=harness_cls.name, account=account, payload=payload, @@ -202,7 +204,7 @@ def make_agent_result(output, *, agent="agent", error=None, cost_usd=None, cost_ ) -def finish_agent_run(call, *, status=None, last_error=None): +async def finish_agent_run(call, *, status=None, last_error=None): # Mark a seeded AgentCall finished (prod builds finished rows via AgentCall.record). from druks.database import db_session from druks.durable.enums import AgentCallStatus @@ -210,32 +212,32 @@ def finish_agent_run(call, *, status=None, last_error=None): call.status = (status or AgentCallStatus.SUCCEEDED).value call.last_error = last_error call.finished_at = Base.utc_now() - db_session().flush() + await db_session().flush() return call -def make_test_note(body: str = "a note"): +async def make_test_note(body: str = "a note"): """The platform suite's subject. It belongs to the proof app, not to ship — platform behavior must hold for any app's rows.""" from druks_field_notes.models import Note - return Note.create(body=body) + return await Note.create(body=body) -def seed_note_run(session, *, note=None, state: str = "running", **kwargs): +async def seed_note_run(session, *, note=None, state: str = "running", **kwargs): """A run on a note, seeding one if the caller has none.""" from druks.testing import seed_run from druks_field_notes.workflows import Summarize - subject = note or make_test_note() - return seed_run(session, kind=Summarize.kind, subject=subject, state=state, **kwargs) + subject = note or await make_test_note() + return await seed_run(session, kind=Summarize.kind, subject=subject, state=state, **kwargs) -def seed_note_agent_run(*, agent: str = "implement", model: str = "gpt-5.5", **kwargs): +async def seed_note_agent_run(*, agent: str = "implement", model: str = "gpt-5.5", **kwargs): """A run on a fresh note with one agent call on it — the call is what the caller wants.""" from druks.database import db_session from druks.testing import seed_call session = db_session() - run = seed_note_run(session, **kwargs) - return seed_call(session, run, agent, status="running", model=model) + run = await seed_note_run(session, **kwargs) + return await seed_call(session, run, agent, status="running", model=model) diff --git a/backend/tests/druks-field_notes/druks_field_notes/models.py b/backend/tests/druks-field_notes/druks_field_notes/models.py index e8423812..3a11e15a 100644 --- a/backend/tests/druks-field_notes/druks_field_notes/models.py +++ b/backend/tests/druks-field_notes/druks_field_notes/models.py @@ -18,33 +18,33 @@ class Note(StoredSubject): created_at: Mapped[datetime] = mapped_column(default=StoredSubject.utc_now) @classmethod - def create(cls, *, body: str) -> "Note": + async def create(cls, *, body: str) -> "Note": session = db_session() note = cls(body=body) session.add(note) - session.flush() + await session.flush() return note @classmethod - def get(cls, note_id: int) -> "Note | None": - return db_session().get(cls, note_id) + async def get(cls, note_id: int) -> "Note | None": + return await db_session().get(cls, note_id) @classmethod - def list_recent(cls, *, limit: int = 100) -> list["Note"]: + async def list_recent(cls, *, limit: int = 100) -> list["Note"]: stmt = select(cls).order_by(cls.created_at.desc(), cls.id.desc()).limit(limit) - return list(db_session().scalars(stmt)) + return list(await db_session().scalars(stmt)) - def save_gist(self, gist: str) -> None: + async def save_gist(self, gist: str) -> None: self.gist = gist - db_session().flush() + await db_session().flush() def get_summary(self) -> NoteSummary: return NoteSummary.model_validate(self) @classmethod - def list_summaries(cls, account_id: str | None) -> list[NoteSummary]: + async def list_summaries(cls, account_id: str | None) -> list[NoteSummary]: # How many the board shows is an operator knob, so it lives on the app. from druks_field_notes.app import FieldNotes - notes = cls.list_recent(limit=FieldNotes.settings().board_size) + notes = await cls.list_recent(limit=(await FieldNotes.settings()).board_size) return [note.get_summary() for note in notes] diff --git a/backend/tests/druks-field_notes/druks_field_notes/routes.py b/backend/tests/druks-field_notes/druks_field_notes/routes.py index c603bdce..0891c0a2 100644 --- a/backend/tests/druks-field_notes/druks_field_notes/routes.py +++ b/backend/tests/druks-field_notes/druks_field_notes/routes.py @@ -12,11 +12,11 @@ @router.get("", response_model=list[NoteSummary], response_model_by_alias=True) async def list_notes() -> list[NoteSummary]: - return [note.get_summary() for note in Note.list_recent()] + return [note.get_summary() for note in await Note.list_recent()] @router.post("", status_code=status.HTTP_201_CREATED) async def write_note(body: Annotated[str, Body(embed=True)]) -> dict[str, int]: - note = Note.create(body=body) + note = await Note.create(body=body) await Summarize.dispatch(note=note) return {"id": note.id} diff --git a/backend/tests/druks-field_notes/druks_field_notes/subscribers.py b/backend/tests/druks-field_notes/druks_field_notes/subscribers.py index 0af60c33..fe74bc2e 100644 --- a/backend/tests/druks-field_notes/druks_field_notes/subscribers.py +++ b/backend/tests/druks-field_notes/druks_field_notes/subscribers.py @@ -10,4 +10,4 @@ async def note_summarized(*, subject: Note, **_: object) -> None: # A finished summarize is a milestone worth its own feed row. The workflow # lifecycle is the trigger; the app only reacts. - FieldNotes.record_event(type="summarized", subject=subject) + await FieldNotes.record_event(type="summarized", subject=subject) diff --git a/backend/tests/druks-field_notes/druks_field_notes/workflows.py b/backend/tests/druks-field_notes/druks_field_notes/workflows.py index 9086fa45..55d31041 100644 --- a/backend/tests/druks-field_notes/druks_field_notes/workflows.py +++ b/backend/tests/druks-field_notes/druks_field_notes/workflows.py @@ -11,11 +11,11 @@ class Summarize(Workflow): subject = Note async def run(self) -> None: - note = self.subject + note = await self.subject # The note body is the agent's prompt context; the gist it returns is the # app's own domain result, saved onto the note. result = await FieldNotes.summarize(note_body=note.body) - note.save_gist(result.gist) + await note.save_gist(result.gist) @classmethod async def dispatch(cls, *, note: Note) -> str: diff --git a/backend/tests/druks-field_notes/tests/test_app.py b/backend/tests/druks-field_notes/tests/test_app.py index b046a64c..f542f520 100644 --- a/backend/tests/druks-field_notes/tests/test_app.py +++ b/backend/tests/druks-field_notes/tests/test_app.py @@ -4,12 +4,12 @@ from pydantic import ValidationError -def test_the_board_honors_the_board_size(druks_db): - Note.create(body="first") - newest = Note.create(body="second") - FieldNotes.override_setting("board_size", 1) +async def test_the_board_honors_the_board_size(druks_db): + await Note.create(body="first") + newest = await Note.create(body="second") + await FieldNotes.override_setting("board_size", 1) - summaries = Note.list_summaries(None) + summaries = await Note.list_summaries(None) assert [summary.id for summary in summaries] == [str(newest.id)] assert summaries[0].body == "second" diff --git a/backend/tests/druks-field_notes/tests/test_models.py b/backend/tests/druks-field_notes/tests/test_models.py index 3c904d31..4b4952a7 100644 --- a/backend/tests/druks-field_notes/tests/test_models.py +++ b/backend/tests/druks-field_notes/tests/test_models.py @@ -1,13 +1,13 @@ from druks_field_notes.models import Note -def test_note_create_list_and_save_gist(druks_db): - first = Note.create(body="the pump ran hot") - second = Note.create(body="the pressure held") +async def test_note_create_list_and_save_gist(druks_db): + first = await Note.create(body="the pump ran hot") + second = await Note.create(body="the pressure held") - assert Note.list_recent(limit=1) == [second] + assert await Note.list_recent(limit=1) == [second] - first.save_gist("The pump ran hot.") + await first.save_gist("The pump ran hot.") - saved = Note.get(first.id) + saved = await Note.get(first.id) assert saved.gist == "The pump ran hot." diff --git a/backend/tests/druks-field_notes/tests/test_routes.py b/backend/tests/druks-field_notes/tests/test_routes.py index ccb17cfc..544466e2 100644 --- a/backend/tests/druks-field_notes/tests/test_routes.py +++ b/backend/tests/druks-field_notes/tests/test_routes.py @@ -2,7 +2,7 @@ from druks_field_notes.workflows import Summarize -def test_notes_routes_create_and_list_notes(druks_client, monkeypatch): +async def test_notes_routes_create_and_list_notes(druks_client, monkeypatch): # The route dispatches from its own session and that session closes with the # request, so read the note's id here while the row is still attached. summarized = [] @@ -13,17 +13,17 @@ async def dispatch(*, note): monkeypatch.setattr(Summarize, "dispatch", staticmethod(dispatch)) - created = druks_client.post( + created = await druks_client.post( "/api/field_notes/notes", json={"body": "the pump ran hot"}, ) assert created.status_code == 201 - note = Note.get(created.json()["id"]) + note = await Note.get(created.json()["id"]) assert note.body == "the pump ran hot" assert summarized == [note.id] - listed = druks_client.get("/api/field_notes/notes") + listed = await druks_client.get("/api/field_notes/notes") assert listed.status_code == 200 assert listed.json()[0]["body"] == "the pump ran hot" diff --git a/backend/tests/druks-field_notes/tests/test_workflows.py b/backend/tests/druks-field_notes/tests/test_workflows.py index 22cce517..001a1b96 100644 --- a/backend/tests/druks-field_notes/tests/test_workflows.py +++ b/backend/tests/druks-field_notes/tests/test_workflows.py @@ -8,14 +8,14 @@ async def test_summarize_writes_the_gist(druks_db, monkeypatch): - note = Note.create(body="the pump ran hot on the second pass") + note = await Note.create(body="the pump ran hot on the second pass") summarize = mock.AsyncMock(return_value=GistOutput(gist="The pump ran hot on the second pass.")) monkeypatch.setattr(FieldNotes, "summarize", staticmethod(summarize)) await run_workflow(Summarize, subject=note) summarize.assert_awaited_once_with(note_body=note.body) - assert Note.get(note.id).gist == "The pump ran hot on the second pass." + assert (await Note.get(note.id)).gist == "The pump ran hot on the second pass." async def test_dispatch_starts_the_workflow_for_the_note(monkeypatch): diff --git a/backend/tests/ship/factories.py b/backend/tests/ship/factories.py index 3fbecad9..bd42da98 100644 --- a/backend/tests/ship/factories.py +++ b/backend/tests/ship/factories.py @@ -4,18 +4,18 @@ from uuid_utils import uuid7 -def make_test_work_item(*, repo: str, **kwargs): +async def make_test_work_item(*, repo: str, **kwargs): """A WorkItem with the Project / ProjectRepo binding it requires. Every item carries a ticket key, unique per source — tests that don't care get one.""" - project = Project.get_for_repo(repo) + project = await Project.get_for_repo(repo) if not project: - project = Project.create(name=repo) - ProjectRepo.create(project_id=project.id, full_name=repo) + project = await Project.create(name=repo) + await ProjectRepo.create(project_id=project.id, full_name=repo) kwargs.setdefault("ticket_key", f"TEST-{uuid7()}") - return WorkItem.create(project_id=project.id, repo=repo, **kwargs) + return await WorkItem.create(project_id=project.id, repo=repo, **kwargs) -def seed_build_run( +async def seed_build_run( session, *, work_item_id: int, @@ -29,8 +29,8 @@ def seed_build_run( timeline; it finds the item through the subject it was started for.""" if state == "parked" and not input_gate: input_gate = "review" # a parked run always has a gate; derivation needs it - item = WorkItem.get(work_item_id) - run = seed_run( + item = await WorkItem.get(work_item_id) + run = await seed_run( session, kind=Build.kind, subject=item, @@ -40,5 +40,5 @@ def seed_build_run( failure=failure, account_id=account_id or "system", ) - session.flush() + await session.flush() return run diff --git a/backend/tests/ship/test_api_work_items.py b/backend/tests/ship/test_api_work_items.py index 631ca6a3..dcda7149 100644 --- a/backend/tests/ship/test_api_work_items.py +++ b/backend/tests/ship/test_api_work_items.py @@ -3,7 +3,7 @@ import pytest from druks.contrib.ship.models import WorkItem -from druks.testing import seed_call +from druks.testing import asgi_client, seed_call from fastapi.testclient import TestClient from ship.factories import make_test_work_item, seed_build_run @@ -16,17 +16,16 @@ } -def _build_client(tmp_path): +def _build_app(tmp_path): from druks.testing import configure_app_for_test, make_settings settings = make_settings(tmp_path) - app = configure_app_for_test(settings=settings) - return TestClient(app) + return configure_app_for_test(settings=settings) @pytest.fixture -def client(tmp_path: Path, druks_db): - with _build_client(tmp_path) as client: +async def client(tmp_path: Path, druks_db): + async with asgi_client(_build_app(tmp_path)) as client: yield client @@ -37,12 +36,12 @@ def client(tmp_path: Path, druks_db): } -def _seed_op(druks_db, work_item_id, *, kind="implement", state, input_gate=None): +async def _seed_op(druks_db, work_item_id, *, kind="implement", state, input_gate=None): """A build run on the item in ``state`` whose latest agent call is ``kind``. When a run already exists, advance it (re-trigger = a fresh round on the same item).""" if state == "running" and input_gate: - run = seed_build_run( + run = await seed_build_run( druks_db, work_item_id=work_item_id, state="parked", @@ -50,16 +49,16 @@ def _seed_op(druks_db, work_item_id, *, kind="implement", state, input_gate=None input_request=_GATE_REQUESTS.get(input_gate), ) else: - run = seed_build_run(druks_db, work_item_id=work_item_id, state=_RUN_STATE[state]) - seed_call(druks_db, run, kind) + run = await seed_build_run(druks_db, work_item_id=work_item_id, state=_RUN_STATE[state]) + await seed_call(druks_db, run, kind) -def _resolve(repo, pr_number, *, merged=True): +async def _resolve(repo, pr_number, *, merged=True): """GitHub's verdict on a work item's PR — what lands it in History, as the merge handler stores it.""" - item = WorkItem.get_for_pr(repo=repo, pr_number=pr_number) + item = await WorkItem.get_for_pr(repo=repo, pr_number=pr_number) if item: - item.resolve(merged=merged, at=datetime.now(UTC)) + await item.resolve(merged=merged, at=datetime.now(UTC)) # The generic subject read-side — Build declares subject = WorkItem, so the @@ -68,42 +67,44 @@ def _resolve(repo, pr_number, *, merged=True): # platform's. See test_generic_subjects.py for the platform-side contract. -def test_subject_list_shows_active_and_excludes_resolved(client: TestClient, druks_db): +async def test_subject_list_shows_active_and_excludes_resolved(client: TestClient, druks_db): repo = "ClawHaven/acme-app" - building = make_test_work_item(title="building", repo=repo).id - _seed_op(druks_db, building, state="running") + building = (await make_test_work_item(title="building", repo=repo)).id + await _seed_op(druks_db, building, state="running") # Merged → History, not the active board. - done = make_test_work_item(title="merged one", repo=repo).id - WorkItem.get(done).update(pr_number=1) - _seed_op(druks_db, done, state="finished") - _resolve(repo, 1) - - rows = {r["summary"]["title"]: r for r in client.get("/api/ship/work_item").json()["rows"]} + done = (await make_test_work_item(title="merged one", repo=repo)).id + await (await WorkItem.get(done)).update(pr_number=1) + await _seed_op(druks_db, done, state="finished") + await _resolve(repo, 1) + + rows = { + r["summary"]["title"]: r for r in (await client.get("/api/ship/work_item")).json()["rows"] + } assert "building" in rows assert "merged one" not in rows assert rows["building"]["status"]["state"] == "running" assert rows["building"]["summary"]["resolution"] is None -def test_subject_detail_composes_summary_status_and_timeline(client: TestClient, druks_db): - item = make_test_work_item( +async def test_subject_detail_composes_summary_status_and_timeline(client: TestClient, druks_db): + item = await make_test_work_item( title="detail", repo="ClawHaven/acme-app", source="linear", ticket_key="ACME-5", ticket_url="https://linear.app/acme/issue/ACME-5/detail", ) - WorkItem.get(item.id).update(pr_number=8) - run = seed_build_run( + await (await WorkItem.get(item.id)).update(pr_number=8) + run = await seed_build_run( druks_db, work_item_id=item.id, state="parked", input_gate="review_plan", input_request={"next_action": "approve_plan", "label": "Approve plan"}, ) - seed_call(druks_db, run, "generate_plan") + await seed_call(druks_db, run, "generate_plan") - detail = client.get(f"/api/ship/work_item/{item.id}").json() + detail = (await client.get(f"/api/ship/work_item/{item.id}")).json() summary = detail["summary"] assert summary["id"] == str(item.id) assert summary["ticketKey"] == "ACME-5" @@ -121,57 +122,57 @@ def test_subject_detail_composes_summary_status_and_timeline(client: TestClient, assert [call["agent"] for call in entry["agentCalls"]] == ["generate_plan"] -def test_subject_detail_unknown_is_404(client: TestClient): - assert client.get("/api/ship/work_item/9999").status_code == 404 +async def test_subject_detail_unknown_is_404(client: TestClient): + assert (await client.get("/api/ship/work_item/9999")).status_code == 404 -def test_pending_gate_surfaces_input_request_on_the_run(druks_db): +async def test_pending_gate_surfaces_input_request_on_the_run(druks_db): # A gate is run-level: the parked run carries its own ask on the timeline, # with its agent calls in execution order underneath. - item = make_test_work_item(repo="ClawHaven/acme-app", title="x") - run = seed_build_run( + item = await make_test_work_item(repo="ClawHaven/acme-app", title="x") + run = await seed_build_run( druks_db, work_item_id=item.id, state="parked", input_gate="review_plan", input_request={"next_action": "approve_plan", "label": "Approve plan"}, ) - seed_call(druks_db, run, "generate_plan") - seed_call(druks_db, run, "review_plan") + await seed_call(druks_db, run, "generate_plan") + await seed_call(druks_db, run, "review_plan") - (entry,) = item.get_timeline() + (entry,) = await item.get_timeline() assert entry.input_request == {"next_action": "approve_plan", "label": "Approve plan"} assert entry.state == "parked" assert [call.agent for call in entry.agent_calls] == ["generate_plan", "review_plan"] -def test_detail_surfaces_running_run_before_its_first_call(druks_db): +async def test_detail_surfaces_running_run_before_its_first_call(druks_db): """The detail timeline surfaces a run that is running before its first agent call exists — the sandbox spin-up window the operator needs to see.""" - item = make_test_work_item(repo="ClawHaven/acme-app", title="x") - seed_build_run(druks_db, work_item_id=item.id, state="running") + item = await make_test_work_item(repo="ClawHaven/acme-app", title="x") + await seed_build_run(druks_db, work_item_id=item.id, state="running") - (entry,) = item.get_timeline() + (entry,) = await item.get_timeline() assert entry.state == "running" assert entry.agent_calls == [] # surfaces even with no call yet -def test_history_returns_only_done_work_items(client: TestClient, druks_db): +async def test_history_returns_only_done_work_items(client: TestClient, druks_db): repo = "ClawHaven/acme-app" # Merged → history. - done_id = make_test_work_item(title="merged one", repo=repo).id - WorkItem.get(done_id).update(pr_number=1) - _seed_op(druks_db, done_id, state="finished") - _resolve(repo, 1) + done_id = (await make_test_work_item(title="merged one", repo=repo)).id + await (await WorkItem.get(done_id)).update(pr_number=1) + await _seed_op(druks_db, done_id, state="finished") + await _resolve(repo, 1) # Running → active. - running_id = make_test_work_item(title="still running", repo=repo).id - _seed_op(druks_db, running_id, state="running") + running_id = (await make_test_work_item(title="still running", repo=repo)).id + await _seed_op(druks_db, running_id, state="running") # Failed (no merge) → active "needs you", NOT history (the whole point). - failed_id = make_test_work_item(title="broke", repo=repo).id - WorkItem.get(failed_id).update(pr_number=2) - _seed_op(druks_db, failed_id, state="failed") + failed_id = (await make_test_work_item(title="broke", repo=repo)).id + await (await WorkItem.get(failed_id)).update(pr_number=2) + await _seed_op(druks_db, failed_id, state="failed") - items = client.get("/api/ship/work-items/history").json()["items"] + items = (await client.get("/api/ship/work-items/history")).json()["items"] titles = [it["title"] for it in items] assert "merged one" in titles assert "still running" not in titles @@ -180,62 +181,62 @@ def test_history_returns_only_done_work_items(client: TestClient, druks_db): assert resolved["resolution"] == "merged" -def test_pr_closed_without_merge_is_closed_in_history(client: TestClient, druks_db): +async def test_pr_closed_without_merge_is_closed_in_history(client: TestClient, druks_db): repo = "ClawHaven/acme-app" # A build parked on the operator, whose PR was then closed without merging. - wid = make_test_work_item(title="abandoned", repo=repo).id - WorkItem.get(wid).update(pr_number=7) - _seed_op(druks_db, wid, state="finished") - _resolve(repo, 7, merged=False) + wid = (await make_test_work_item(title="abandoned", repo=repo)).id + await (await WorkItem.get(wid)).update(pr_number=7) + await _seed_op(druks_db, wid, state="finished") + await _resolve(repo, 7, merged=False) - items = client.get("/api/ship/work-items/history").json()["items"] + items = (await client.get("/api/ship/work-items/history")).json()["items"] row = next(it for it in items if it["title"] == "abandoned") assert row["resolution"] == "closed" # History's time column is the verdict's, not the row's last touch. - assert datetime.fromisoformat(row["updatedAt"]) == WorkItem.get(wid).resolved_at + assert datetime.fromisoformat(row["updatedAt"]) == (await WorkItem.get(wid)).resolved_at -def test_history_clamps_limit(client: TestClient, druks_db): +async def test_history_clamps_limit(client: TestClient, druks_db): for i in range(3): - wid = make_test_work_item(title=f"merged {i}", repo="ClawHaven/acme-app").id - WorkItem.get(wid).update(pr_number=i + 1) - _seed_op(druks_db, wid, state="finished") - _resolve("ClawHaven/acme-app", i + 1) + wid = (await make_test_work_item(title=f"merged {i}", repo="ClawHaven/acme-app")).id + await (await WorkItem.get(wid)).update(pr_number=i + 1) + await _seed_op(druks_db, wid, state="finished") + await _resolve("ClawHaven/acme-app", i + 1) # limit > cap → clamps down, doesn't 400. - response = client.get("/api/ship/work-items/history?limit=10000") + response = await client.get("/api/ship/work-items/history?limit=10000") assert response.status_code == 200 items = response.json()["items"] assert len(items) == 3 # all three merged; cap doesn't truncate here # limit < 1 → clamps up to 1. - response = client.get("/api/ship/work-items/history?limit=0") + response = await client.get("/api/ship/work-items/history?limit=0") assert response.status_code == 200 items = response.json()["items"] assert len(items) == 1 -def test_repeated_runs_on_one_subject_each_surface_separately(druks_db): +async def test_repeated_runs_on_one_subject_each_surface_separately(druks_db): # The timeline must not collapse repeated runs to only the newest one. - item = make_test_work_item(repo="ClawHaven/acme-app", title="repeated") + item = await make_test_work_item(repo="ClawHaven/acme-app", title="repeated") for _ in range(3): - seed_build_run(druks_db, work_item_id=item.id, state="finished") + await seed_build_run(druks_db, work_item_id=item.id, state="finished") - entries = item.get_timeline() + entries = await item.get_timeline() assert [entry.kind for entry in entries] == ["ship.build"] * 3 assert len({entry.id for entry in entries}) == 3 -def test_timeline_shows_every_build_attempt(druks_db): +async def test_timeline_shows_every_build_attempt(druks_db): # Each build attempt is its own run; the timeline shows them all, with a # failed attempt's failure carried on its run. - item = make_test_work_item(repo="ClawHaven/acme-app", title="x", ticket_key="ACME-1") - run1 = seed_build_run(druks_db, work_item_id=item.id, state="failed", failure="boom") - run2 = seed_build_run(druks_db, work_item_id=item.id, state="failed") - seed_call(druks_db, run1, "generate_plan", status="failed", last_error="boom") - seed_call(druks_db, run2, "generate_plan", status="failed") + item = await make_test_work_item(repo="ClawHaven/acme-app", title="x", ticket_key="ACME-1") + run1 = await seed_build_run(druks_db, work_item_id=item.id, state="failed", failure="boom") + run2 = await seed_build_run(druks_db, work_item_id=item.id, state="failed") + await seed_call(druks_db, run1, "generate_plan", status="failed", last_error="boom") + await seed_call(druks_db, run2, "generate_plan", status="failed") - entries = item.get_timeline() + entries = await item.get_timeline() assert len(entries) == 2 assert all(e.agent_calls[0].agent == "generate_plan" for e in entries) assert any(e.failure == "boom" for e in entries) @@ -246,8 +247,8 @@ async def test_subject_activity_surfaces_running_phase(druks_db, monkeypatch): # surfaces it ("Building sandbox VM…") — finer than the lifecycle status. from druks.contrib.ship import app as ship_app - item = make_test_work_item(repo="ClawHaven/acme-app", title="x") - seed_build_run(druks_db, work_item_id=item.id, state="running") + item = await make_test_work_item(repo="ClawHaven/acme-app", title="x") + await seed_build_run(druks_db, work_item_id=item.id, state="running") async def phase(_run_id): return "provisioning_vm" @@ -263,7 +264,7 @@ async def test_subject_activity_none_when_not_running(druks_db): # A run parked on a gate isn't working — no live sub-phase. from druks.contrib.ship import app as ship_app - item = make_test_work_item(repo="ClawHaven/acme-app", title="x") - seed_build_run(druks_db, work_item_id=item.id, state="parked", input_gate="review_plan") + item = await make_test_work_item(repo="ClawHaven/acme-app", title="x") + await seed_build_run(druks_db, work_item_id=item.id, state="parked", input_gate="review_plan") assert await ship_app.Ship.get_subject_activity(item) is None diff --git a/backend/tests/ship/test_app_config.py b/backend/tests/ship/test_app_config.py index 65673e2a..affdf380 100644 --- a/backend/tests/ship/test_app_config.py +++ b/backend/tests/ship/test_app_config.py @@ -57,7 +57,11 @@ class TestFetchFile: def _wire(self, monkeypatch, tmp_path, github): settings = make_settings(tmp_path) monkeypatch.setattr("druks.apps.fetcher.load_settings", lambda: settings) - monkeypatch.setattr("druks.apps.fetcher.get_github_client", lambda: github) + + async def _client(): + return github + + monkeypatch.setattr("druks.apps.fetcher.get_github_client", _client) async def test_404_is_cached_as_empty(self, monkeypatch, tmp_path): from druks.apps.fetcher import fetch_file @@ -133,10 +137,10 @@ class TestLoadPolicyAndProfile: needed.""" @pytest.fixture(autouse=True) - def _passthrough_step(self, monkeypatch, druks_db): + async def _passthrough_step(self, monkeypatch, druks_db): from druks.durable.engine import configure_engine - configure_engine(druks_db.connection()) + configure_engine(await druks_db.connection()) async def _run_step(_options, fn): return await fn() @@ -163,8 +167,8 @@ async def _live(cls, repo): # A build only runs for a registered repo; seed it so live resolution # reads its (as-yet-empty) profile. - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name=REPO) + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name=REPO) resolved = await self._flow(repo=REPO)._load_policy_and_profile() assert RepoPolicy.model_validate(resolved["policy"]).sandbox.image == "live" diff --git a/backend/tests/ship/test_build_board_membership.py b/backend/tests/ship/test_build_board_membership.py index f5736458..1bd79bf1 100644 --- a/backend/tests/ship/test_build_board_membership.py +++ b/backend/tests/ship/test_build_board_membership.py @@ -7,62 +7,62 @@ from ship.factories import make_test_work_item, seed_build_run -def _board_ids(druks_db): - druks_db.expire_all() - return {row.id for row in WorkItem.list_summaries(None)} +async def _board_ids(druks_db): + druks_db.expunge_all() + return {row.id for row in await WorkItem.list_summaries(None)} -def _resolve(item, *, merged=True, at=None): - item.resolve(merged=merged, at=at or datetime.now(UTC)) +async def _resolve(item, *, merged=True, at=None): + await item.resolve(merged=merged, at=at or datetime.now(UTC)) @pytest.mark.parametrize( "state", ["scheduled", "running", "parked", "failed", "finished", "cancelled"] ) -def test_an_unresolved_item_holds_the_board(druks_db, state): - item = make_test_work_item(repo="ClawHaven/acme-app", title=f"build {state}") - seed_build_run(druks_db, work_item_id=item.id, state=state) - assert str(item.id) in _board_ids(druks_db) +async def test_an_unresolved_item_holds_the_board(druks_db, state): + item = await make_test_work_item(repo="ClawHaven/acme-app", title=f"build {state}") + await seed_build_run(druks_db, work_item_id=item.id, state=state) + assert str(item.id) in await _board_ids(druks_db) @pytest.mark.parametrize("state", ["running", "failed", "finished"]) -def test_a_resolved_pr_leaves_the_board_whatever_its_run_says(druks_db, state): - item = make_test_work_item(repo="ClawHaven/acme-app", title=f"resolved {state}") - seed_build_run(druks_db, work_item_id=item.id, state=state) - _resolve(item) - assert str(item.id) not in _board_ids(druks_db) +async def test_a_resolved_pr_leaves_the_board_whatever_its_run_says(druks_db, state): + item = await make_test_work_item(repo="ClawHaven/acme-app", title=f"resolved {state}") + await seed_build_run(druks_db, work_item_id=item.id, state=state) + await _resolve(item) + assert str(item.id) not in await _board_ids(druks_db) -def test_a_redispatched_item_returns_to_the_board(druks_db): - item = make_test_work_item(repo="ClawHaven/acme-app", title="rebuilt") - _resolve(item, merged=False) +async def test_a_redispatched_item_returns_to_the_board(druks_db): + item = await make_test_work_item(repo="ClawHaven/acme-app", title="rebuilt") + await _resolve(item, merged=False) - item.start_attempt() + await item.start_attempt() - assert str(item.id) in _board_ids(druks_db) + assert str(item.id) in await _board_ids(druks_db) -def test_the_board_does_not_ask_about_runs(druks_db): - item = make_test_work_item(repo="ClawHaven/acme-app", title="never dispatched") - assert str(item.id) in _board_ids(druks_db) +async def test_the_board_does_not_ask_about_runs(druks_db): + item = await make_test_work_item(repo="ClawHaven/acme-app", title="never dispatched") + assert str(item.id) in await _board_ids(druks_db) -def test_history_holds_the_resolved_newest_verdict_first(druks_db): - older = make_test_work_item(repo="ClawHaven/acme-app", title="closed first") - _resolve(older, merged=False, at=datetime.now(UTC) - timedelta(minutes=1)) - newer = make_test_work_item(repo="ClawHaven/acme-app", title="merged after") - _resolve(newer) - make_test_work_item(repo="ClawHaven/acme-app", title="still open") +async def test_history_holds_the_resolved_newest_verdict_first(druks_db): + older = await make_test_work_item(repo="ClawHaven/acme-app", title="closed first") + await _resolve(older, merged=False, at=datetime.now(UTC) - timedelta(minutes=1)) + newer = await make_test_work_item(repo="ClawHaven/acme-app", title="merged after") + await _resolve(newer) + await make_test_work_item(repo="ClawHaven/acme-app", title="still open") - assert [item.id for item in WorkItem.list_handoff()] == [newer.id, older.id] + assert [item.id for item in await WorkItem.list_handoff()] == [newer.id, older.id] -def test_the_newest_run_speaks_for_the_item(druks_db): +async def test_the_newest_run_speaks_for_the_item(druks_db): # One rule, two implementations — the bulk board query and the per-subject # status must name the same driving run or the board and its lanes disagree. - item = make_test_work_item(repo="ClawHaven/acme-app", title="two runs") - seed_build_run(druks_db, work_item_id=item.id, state="cancelled") - seed_build_run(druks_db, work_item_id=item.id, state="parked", input_gate="review") - druks_db.expire_all() - assert item.get_status().state == "parked" - assert str(item.id) in _board_ids(druks_db) + item = await make_test_work_item(repo="ClawHaven/acme-app", title="two runs") + await seed_build_run(druks_db, work_item_id=item.id, state="cancelled") + await seed_build_run(druks_db, work_item_id=item.id, state="parked", input_gate="review") + druks_db.expunge_all() + assert (await item.get_status()).state == "parked" + assert str(item.id) in await _board_ids(druks_db) diff --git a/backend/tests/ship/test_build_dispatch.py b/backend/tests/ship/test_build_dispatch.py index fddc81b6..3f3ca184 100644 --- a/backend/tests/ship/test_build_dispatch.py +++ b/backend/tests/ship/test_build_dispatch.py @@ -8,8 +8,8 @@ from ship.factories import make_test_work_item -def _connect_github() -> None: - ServiceIdentity.connect( +async def _connect_github() -> None: + await ServiceIdentity.connect( "github", identity={"app_id": "1", "slug": "druks-operator"}, secrets={"private_key": "operator-pem", "webhook_secret": "hook-secret"}, @@ -36,12 +36,12 @@ async def test_dispatch_leaves_the_item_alone(druks_db, monkeypatch) -> None: """Dispatch starts the build and touches nothing else — clearing the previous attempt is the scheduled reaction's (test_lane_reactions), and a duplicate dispatch never makes that announcement.""" - _connect_github() - seed_run(druks_db, kind=Build.kind, run_id="run-old") - seed_run(druks_db, kind=Build.kind, run_id="run-new") - item = make_test_work_item(repo="o/r", title="t", ticket_key="ACME-2") - item.update(pr_number=7, branch="agent/old") - item.resolve(merged=False, at=datetime.now(UTC)) + await _connect_github() + await seed_run(druks_db, kind=Build.kind, run_id="run-old") + await seed_run(druks_db, kind=Build.kind, run_id="run-new") + item = await make_test_work_item(repo="o/r", title="t", ticket_key="ACME-2") + await item.update(pr_number=7, branch="agent/old") + await item.resolve(merged=False, at=datetime.now(UTC)) async def fake_start(cls, **kwargs): return "run-new" @@ -60,7 +60,7 @@ async def test_dispatch_stands_down_without_github_instead_of_raising( """The tracker delivery already succeeded — a raise here would 5xx the webhook into provider redelivery. No identity: log the not-connected direction and start nothing.""" - item = make_test_work_item(repo="o/r", title="t", ticket_key="ACME-8") + item = await make_test_work_item(repo="o/r", title="t", ticket_key="ACME-8") started = [] async def fake_start(cls, **kwargs): @@ -84,8 +84,10 @@ async def test_the_tracker_funnel_swallows_the_missing_identity(druks_db, monkey from druks.contrib.ship import subscribers # noqa: F401 — the import registers it from druks.contrib.ship.app import Ship - settings = Ship.settings() - item = make_test_work_item(repo="o/r", title="t", ticket_key="ACME-9", source=settings.tracker) + settings = await Ship.settings() + item = await make_test_work_item( + repo="o/r", title="t", ticket_key="ACME-9", source=settings.tracker + ) started = [] async def fake_start(cls, **kwargs): @@ -107,9 +109,9 @@ async def test_dispatch_merged_noop_still_precedes_the_identity_guard( ) -> None: """A merged item's redelivery keeps its own no-op — the identity guard only decides deliveries that would otherwise start.""" - item = make_test_work_item(repo="o/r", title="t", ticket_key="ACME-10") - item.update(pr_number=7, branch="agent/old") - item.resolve(merged=True, at=datetime.now(UTC)) + item = await make_test_work_item(repo="o/r", title="t", ticket_key="ACME-10") + await item.update(pr_number=7, branch="agent/old") + await item.resolve(merged=True, at=datetime.now(UTC)) async def fake_start(cls, **kwargs): raise AssertionError("a merged item never starts") @@ -152,13 +154,13 @@ async def fake_start(cls, **kwargs): assert any("no routable repo" in record.getMessage() for record in caplog.records) -def test_update_clears_nullable_with_none_and_skips_omitted(druks_db) -> None: +async def test_update_clears_nullable_with_none_and_skips_omitted(druks_db) -> None: """update() tells a clear from a skip: pr_number=None clears the column, while leaving branch out preserves it.""" - item = make_test_work_item(repo="o/r", title="t", ticket_key="ACME-4") - item.update(pr_number=9, branch="agent/keep") + item = await make_test_work_item(repo="o/r", title="t", ticket_key="ACME-4") + await item.update(pr_number=9, branch="agent/keep") - item.update(pr_number=None) + await item.update(pr_number=None) assert item.pr_number is None assert item.branch == "agent/keep" diff --git a/backend/tests/ship/test_build_plan_phase.py b/backend/tests/ship/test_build_plan_phase.py index 16d5d014..9c092512 100644 --- a/backend/tests/ship/test_build_plan_phase.py +++ b/backend/tests/ship/test_build_plan_phase.py @@ -586,7 +586,7 @@ async def fake_approved_work(): await workflow._work_gate() run = Run(input_request=input_requests[0]) - assert run.get_rendered_ask()["body"] + assert (await run.get_rendered_ask())["body"] async def test_needs_clarification_delivery_stops_the_run(monkeypatch): diff --git a/backend/tests/ship/test_build_prompts.py b/backend/tests/ship/test_build_prompts.py index bf3321f3..c52097bb 100644 --- a/backend/tests/ship/test_build_prompts.py +++ b/backend/tests/ship/test_build_prompts.py @@ -196,16 +196,16 @@ def test_build_prompt_context_covers_template_attrs(): assert not missing, f"BuildPromptContext missing template attrs: {missing}" -def test_get_for_repo_returns_the_repo(druks_db): - project = Project.create(name="acme/widget") - ProjectRepo.create(project_id=project.id, full_name="acme/widget") - assert ProjectRepo.get_for_repo("acme/widget").full_name == "acme/widget" +async def test_get_for_repo_returns_the_repo(druks_db): + project = await Project.create(name="acme/widget") + await ProjectRepo.create(project_id=project.id, full_name="acme/widget") + assert (await ProjectRepo.get_for_repo("acme/widget")).full_name == "acme/widget" -def test_get_for_repo_raises_when_the_repo_was_transferred(druks_db): +async def test_get_for_repo_raises_when_the_repo_was_transferred(druks_db): # Registered under the new name; a run still holding the old name must fail # with the reason, not an opaque NoneType crash. - project = Project.create(name="czpython/druks") - ProjectRepo.create(project_id=project.id, full_name="czpython/druks") + project = await Project.create(name="czpython/druks") + await ProjectRepo.create(project_id=project.id, full_name="czpython/druks") with pytest.raises(FatalError, match="renamed or transferred"): - ProjectRepo.get_for_repo("clawhaven/druks", raise_on_missing=True) + await ProjectRepo.get_for_repo("clawhaven/druks", raise_on_missing=True) diff --git a/backend/tests/ship/test_build_workspace.py b/backend/tests/ship/test_build_workspace.py index 4df55f29..db2e26d2 100644 --- a/backend/tests/ship/test_build_workspace.py +++ b/backend/tests/ship/test_build_workspace.py @@ -81,11 +81,15 @@ async def fake_exec(self: Any, argv: list[str], **_kw: Any) -> Any: monkeypatch.setattr(host_mod.Sandbox, "write_secret", _noop) monkeypatch.setattr(host_mod.Sandbox, "exec", fake_exec) - monkeypatch.setattr( - "druks.contrib.ship.workflows.get_github_client", - lambda: SimpleNamespace(token_for_repo=_token), - ) - monkeypatch.setattr("druks.contrib.ship.workflows.get_review_actor", review_actor) + + async def _github_client(): + return SimpleNamespace(token_for_repo=_token) + + async def _review_actor(): + return review_actor() + + monkeypatch.setattr("druks.contrib.ship.workflows.get_github_client", _github_client) + monkeypatch.setattr("druks.contrib.ship.workflows.get_review_actor", _review_actor) monkeypatch.setattr("druks.sandbox.repo.ensure", fake_ensure) return ensured, execs @@ -164,17 +168,16 @@ def _dispatched_by(monkeypatch: pytest.MonkeyPatch, username: str | None) -> Non async def _bot_git_author() -> tuple[str, str]: return "app[bot]", "1+app[bot]@users.noreply.github.com" - monkeypatch.setattr( - workspace_mod, - "get_github_client", - lambda: SimpleNamespace(get_bot_git_author=_bot_git_author), - ) + async def _client(): + return SimpleNamespace(get_bot_git_author=_bot_git_author) + + monkeypatch.setattr(workspace_mod, "get_github_client", _client) account = SimpleNamespace(username=username) if username else None - monkeypatch.setattr( - workspace_mod, - "Account", - SimpleNamespace(get=lambda _id, *, exclude_system: account), - ) + + async def _get_account(_id, *, exclude_system): + return account + + monkeypatch.setattr(workspace_mod, "Account", SimpleNamespace(get=_get_account)) async def test_set_git_identity_stamps_the_workspace_repo( diff --git a/backend/tests/ship/test_lane_reactions.py b/backend/tests/ship/test_lane_reactions.py index 9124f192..351aced8 100644 --- a/backend/tests/ship/test_lane_reactions.py +++ b/backend/tests/ship/test_lane_reactions.py @@ -20,19 +20,23 @@ async def test_new_build_claims_the_item(druks_db): # A resume replays the body without passing through start(), so this topic — # not RUNNING — is the one that fires once per attempt. - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-1") - item.update(pr_number=7, branch="agent/old") - item.resolve(merged=False, at=datetime.now(UTC)) + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-1" + ) + await item.update(pr_number=7, branch="agent/old") + await item.resolve(merged=False, at=datetime.now(UTC)) await publish(WorkflowEvent.SCHEDULED, subject=item.identity, kind=Build.kind) - refreshed = WorkItem.get(item.id) + refreshed = await WorkItem.get(item.id) assert (refreshed.branch, refreshed.pr_number) == (None, None) assert (refreshed.resolution, refreshed.resolved_at) == (None, None) async def test_cancelled_build_settles_the_item(druks_db): - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-2") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-2" + ) await publish( WorkflowEvent.CANCELLED, @@ -41,18 +45,20 @@ async def test_cancelled_build_settles_the_item(druks_db): failure="operator cancelled", ) - refreshed = WorkItem.get(item.id) + refreshed = await WorkItem.get(item.id) assert refreshed.resolution == "closed" assert refreshed.resolved_at - druks_db.expire_all() - assert str(item.id) not in {summary.id for summary in WorkItem.list_summaries(None)} + druks_db.expunge_all() + assert str(item.id) not in {summary.id for summary in await WorkItem.list_summaries(None)} @pytest.mark.parametrize(("merged", "resolution"), [(True, "merged"), (False, "closed")]) async def test_cancelled_build_preserves_an_existing_resolution(druks_db, merged, resolution): - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-3") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-3" + ) resolved_at = datetime(2026, 8, 8, 12, tzinfo=UTC) - item.resolve(merged=merged, at=resolved_at) + await item.resolve(merged=merged, at=resolved_at) await publish( WorkflowEvent.CANCELLED, @@ -61,13 +67,15 @@ async def test_cancelled_build_preserves_an_existing_resolution(druks_db, merged failure="operator cancelled", ) - refreshed = WorkItem.get(item.id) + refreshed = await WorkItem.get(item.id) assert refreshed.resolution == resolution assert refreshed.resolved_at == resolved_at async def test_failed_build_remains_unresolved_on_the_board(druks_db): - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-4") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-4" + ) await publish( WorkflowEvent.FAILED, @@ -76,10 +84,10 @@ async def test_failed_build_remains_unresolved_on_the_board(druks_db): failure="build failed", ) - refreshed = WorkItem.get(item.id) + refreshed = await WorkItem.get(item.id) assert (refreshed.resolution, refreshed.resolved_at) == (None, None) - druks_db.expire_all() - assert str(item.id) in {summary.id for summary in WorkItem.list_summaries(None)} + druks_db.expunge_all() + assert str(item.id) in {summary.id for summary in await WorkItem.list_summaries(None)} async def test_build_lifecycle_reaches_the_tracker(druks_db, monkeypatch): @@ -89,7 +97,9 @@ async def _push(self, status): pushed.append(status) monkeypatch.setattr(WorkItem, "set_ticket_status", _push) - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-7") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-7" + ) subject = item.identity await publish(WorkflowEvent.RUNNING, subject=subject, kind=Build.kind) @@ -102,8 +112,10 @@ async def _push(self, status): async def test_pr_review_answers_through_the_review_gate(druks_db, monkeypatch): - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-9") - item.update(pr_number=12, branch="agent/acme-9") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-9" + ) + await item.update(pr_number=12, branch="agent/acme-9") run = Run( id=str(uuid7()), kind=Build.kind, @@ -111,8 +123,8 @@ async def test_pr_review_answers_through_the_review_gate(druks_db, monkeypatch): input_request={"presentation": "external", "label": "Review implementation"}, ) druks_db.add(run) - druks_db.flush() - seed_dbos_status( + await druks_db.flush() + await seed_dbos_status( druks_db, run.id, "parked", @@ -150,7 +162,9 @@ async def answer(subject, **reply): async def test_pr_open_reaches_the_work_item(druks_db): - item = make_test_work_item(repo="acme/widget", title="t", source="linear", ticket_key="ACME-8") + item = await make_test_work_item( + repo="acme/widget", title="t", source="linear", ticket_key="ACME-8" + ) await publish( "pr.opened", @@ -160,5 +174,5 @@ async def test_pr_open_reaches_the_work_item(druks_db): branch="agent/eng-8", ) - refreshed = WorkItem.get(item.id) + refreshed = await WorkItem.get(item.id) assert refreshed.pr_number == 12 and refreshed.branch == "agent/eng-8" diff --git a/backend/tests/ship/test_profiling.py b/backend/tests/ship/test_profiling.py index ae485741..b5c62360 100644 --- a/backend/tests/ship/test_profiling.py +++ b/backend/tests/ship/test_profiling.py @@ -11,10 +11,10 @@ @pytest.fixture(autouse=True) -def _passthrough_step(monkeypatch, druks_db): +async def _passthrough_step(monkeypatch, druks_db): # run() is itself a durable step (single-operation workflow) — route it # straight through so the test needs no live DBOS runtime. - configure_engine(druks_db.connection()) + configure_engine(await druks_db.connection()) async def _run_step(_options, fn): return await fn() @@ -24,13 +24,13 @@ async def _run_step(_options, fn): configure_engine(None) -def _seed_repo() -> ProjectRepo: - project = Project.create(name="Acme") - return ProjectRepo.create(project_id=project.id, full_name="acme/widget") +async def _seed_repo() -> ProjectRepo: + project = await Project.create(name="Acme") + return await ProjectRepo.create(project_id=project.id, full_name="acme/widget") -def _seed_skills(*names: str, disabled: tuple[str, ...] = ()) -> None: - collection = SkillCollection.create( +async def _seed_skills(*names: str, disabled: tuple[str, ...] = ()) -> None: + collection = await SkillCollection.create( source="test", name="test skills", skills=[ @@ -66,12 +66,12 @@ async def _no_policy(repo): @pytest.mark.parametrize("refresh_only", [False, True]) async def test_dispatch_shapes_the_profile_start(druks_db, monkeypatch, refresh_only): - ServiceIdentity.connect( + await ServiceIdentity.connect( "github", identity={"app_id": "1", "slug": "druks-operator"}, secrets={"private_key": "operator-pem", "webhook_secret": "hook-secret"}, ) - repo = _seed_repo() + repo = await _seed_repo() calls: list[dict] = [] async def _start(cls, **kwargs): @@ -87,7 +87,7 @@ async def _start(cls, **kwargs): async def test_dispatch_refuses_before_start_without_github(druks_db, monkeypatch): - repo = _seed_repo() + repo = await _seed_repo() async def _start(cls, **kwargs): raise AssertionError("start must not be reached without a GitHub identity") @@ -100,8 +100,8 @@ async def _start(cls, **kwargs): class TestProfileRun: async def test_persists_baseline_and_effective(self, druks_db, monkeypatch): - _seed_skills("django-patterns") - repo = _seed_repo() + await _seed_skills("django-patterns") + repo = await _seed_repo() async def _profiler(*, repo: str): return _profiled() @@ -112,7 +112,7 @@ async def _profiler(*, repo: str): await Profile().run(repo_id=repo.id) # The step commits on its own Session; re-fetch instead of trusting # the identity-mapped `repo` object across that boundary. - repo = ProjectRepo.get(repo.id) + repo = await ProjectRepo.get(repo.id) assert repo.profile["baseline"]["languages"] == ["python"] assert repo.effective_profile["verification"]["lint_commands"] == [ @@ -120,8 +120,8 @@ async def _profiler(*, repo: str): ] async def test_drops_skills_that_are_not_enabled(self, druks_db, monkeypatch): - _seed_skills("django-patterns", "retired-skill", disabled=("retired-skill",)) - repo = _seed_repo() + await _seed_skills("django-patterns", "retired-skill", disabled=("retired-skill",)) + repo = await _seed_repo() async def _profiler(*, repo: str): # The agent picked a disabled skill and one that was never real. @@ -133,12 +133,12 @@ async def _profiler(*, repo: str): monkeypatch.setattr(RepoPolicy, "resolve", staticmethod(_no_policy)) await Profile().run(repo_id=repo.id) - repo = ProjectRepo.get(repo.id) + repo = await ProjectRepo.get(repo.id) assert repo.profile["baseline"]["recommended_skills"] == ["django-patterns"] async def test_pinned_verification_replaces_the_detected_one(self, druks_db, monkeypatch): - repo = _seed_repo() + repo = await _seed_repo() async def _profiler(*, repo: str): return _profiled() @@ -150,7 +150,7 @@ async def _pinning_policy(repo): monkeypatch.setattr(RepoPolicy, "resolve", staticmethod(_pinning_policy)) await Profile().run(repo_id=repo.id) - repo = ProjectRepo.get(repo.id) + repo = await ProjectRepo.get(repo.id) # The pin replaces the whole verification section on the effective profile... assert repo.effective_profile["verification"]["test_commands"] == [ @@ -163,7 +163,7 @@ async def _pinning_policy(repo): ] async def test_pinned_command_keeps_the_check_detected_for_it(self, druks_db, monkeypatch): - repo = _seed_repo() + repo = await _seed_repo() async def _profiler(*, repo: str): return _profiled() @@ -177,7 +177,7 @@ async def _pinning_policy(repo): monkeypatch.setattr(RepoPolicy, "resolve", staticmethod(_pinning_policy)) await Profile().run(repo_id=repo.id) - repo = ProjectRepo.get(repo.id) + repo = await ProjectRepo.get(repo.id) assert repo.effective_profile["verification"]["test_commands"] == [ {"command": "pytest", "ci_check": "Backend / tests"}, @@ -187,9 +187,9 @@ async def _pinning_policy(repo): class TestRefreshOnly: async def test_skips_the_agent_and_reapplies_the_pin(self, druks_db, monkeypatch): - repo = _seed_repo() + repo = await _seed_repo() baseline = _profiled() - repo.set_profile(baseline=baseline, effective=baseline) + await repo.set_profile(baseline=baseline, effective=baseline) async def _boom(*, repo: str): raise AssertionError("refresh_only must not call the repo profiler") @@ -201,7 +201,7 @@ async def _pinning_policy(repo): monkeypatch.setattr(RepoPolicy, "resolve", staticmethod(_pinning_policy)) await Profile().run(repo_id=repo.id, refresh_only=True) - repo = ProjectRepo.get(repo.id) + repo = await ProjectRepo.get(repo.id) # Baseline untouched — only the pin re-applies. assert repo.profile["baseline"]["verification"]["test_commands"] == [ diff --git a/backend/tests/ship/test_project_repo_routes.py b/backend/tests/ship/test_project_repo_routes.py index c8f0d96a..9a708c6e 100644 --- a/backend/tests/ship/test_project_repo_routes.py +++ b/backend/tests/ship/test_project_repo_routes.py @@ -1,16 +1,17 @@ from pathlib import Path import pytest +from druks.database import db_session from fastapi.testclient import TestClient @pytest.fixture -def client(tmp_path: Path, druks_db, monkeypatch): - from druks.testing import configure_app_for_test, make_settings +async def client(tmp_path: Path, druks_db, monkeypatch): + from druks.testing import asgi_client, configure_app_for_test, make_settings monkeypatch.setenv("DRUKS_DATA_DIR", str(tmp_path)) app = configure_app_for_test(settings=make_settings(tmp_path)) - with TestClient(app) as client: + async with asgi_client(app) as client: yield client @@ -29,13 +30,24 @@ async def _dispatch(cls, repo, *, refresh_only=False): return calls -def test_adding_a_repo_dispatches_a_profile_run(client: TestClient, monkeypatch): +async def test_get_project_returns_the_summary_or_404(client: TestClient): + created = (await client.post("/api/ship/projects", json={"name": "Acme"})).json() + + fetched = await client.get(f"/api/ship/projects/{created['id']}") + assert fetched.status_code == 200 + assert fetched.json() == created + assert (await client.get("/api/ship/projects/999999")).status_code == 404 + + +async def test_adding_a_repo_dispatches_a_profile_run(client: TestClient, monkeypatch): calls = _stub_profile_dispatch(monkeypatch) - project = client.post("/api/ship/projects", json={"name": "Acme"}).json() - repo = client.post( - f"/api/ship/projects/{project['id']}/repos", - json={"fullName": "acme/widget"}, + project = (await client.post("/api/ship/projects", json={"name": "Acme"})).json() + repo = ( + await client.post( + f"/api/ship/projects/{project['id']}/repos", + json={"fullName": "acme/widget"}, + ) ).json() assert calls == [ @@ -47,11 +59,11 @@ def test_adding_a_repo_dispatches_a_profile_run(client: TestClient, monkeypatch) assert repo["profile"] == {} -def test_adding_a_repo_survives_when_github_is_not_connected(client: TestClient): +async def test_adding_a_repo_survives_when_github_is_not_connected(client: TestClient): # Registering a repo is metadata; a missing GitHub identity defers profiling # but must not discard the repo — a rollback would 500 and lose it. - project = client.post("/api/ship/projects", json={"name": "Acme"}).json() - response = client.post( + project = (await client.post("/api/ship/projects", json={"name": "Acme"})).json() + response = await client.post( f"/api/ship/projects/{project['id']}/repos", json={"fullName": "acme/widget"}, ) @@ -60,16 +72,16 @@ def test_adding_a_repo_survives_when_github_is_not_connected(client: TestClient) assert response.json()["fullName"] == "acme/widget" -def test_profile_endpoint_dispatches(client: TestClient, monkeypatch): +async def test_profile_endpoint_dispatches(client: TestClient, monkeypatch): # Concurrency is the Profile workflow's subject-unique lock, not the route's # job — the route always dispatches and start() dedups against a live run. from druks.contrib.ship.models import Project, ProjectRepo calls = _stub_profile_dispatch(monkeypatch) - project = Project.create(name="Acme") - repo = ProjectRepo.create(project_id=project.id, full_name="acme/widget") + project = await Project.create(name="Acme") + repo = await ProjectRepo.create(project_id=project.id, full_name="acme/widget") - response = client.post(f"/api/ship/projects/{project.id}/repos/{repo.id}/profile") + response = await client.post(f"/api/ship/projects/{project.id}/repos/{repo.id}/profile") assert response.status_code == 200 assert calls == [ @@ -80,39 +92,39 @@ def test_profile_endpoint_dispatches(client: TestClient, monkeypatch): ] -def test_nested_repo_routes_are_scoped_to_their_project(client: TestClient, monkeypatch): +async def test_nested_repo_routes_are_scoped_to_their_project(client: TestClient, monkeypatch): """PATCH / profile / DELETE reached through the wrong project's URL are 404 and side-effect-free — the routes scope by (project_id, repo_id), not repo_id alone.""" from druks.contrib.ship.models import Project, ProjectRepo profile_calls = _stub_profile_dispatch(monkeypatch) - owner = Project.create(name="Owner") - other = Project.create(name="Other") - repo_id = ProjectRepo.create(project_id=owner.id, full_name="acme/widget").id + owner = await Project.create(name="Owner") + other = await Project.create(name="Other") + repo_id = (await ProjectRepo.create(project_id=owner.id, full_name="acme/widget")).id wrong = f"/api/ship/projects/{other.id}/repos/{repo_id}" - assert client.patch(wrong, json={"purpose": "infra"}).status_code == 404 - assert client.post(f"{wrong}/profile").status_code == 404 - assert client.delete(wrong).status_code == 404 + assert (await client.patch(wrong, json={"purpose": "infra"})).status_code == 404 + assert (await client.post(f"{wrong}/profile")).status_code == 404 + assert (await client.delete(wrong)).status_code == 404 # None of the wrong-parent calls mutated the repo or dispatched a profile run. - assert ProjectRepo.get(repo_id).purpose is None + assert (await ProjectRepo.get(repo_id)).purpose is None assert profile_calls == [] # Through its own project the repo mutates and deletes as normal. right = f"/api/ship/projects/{owner.id}/repos/{repo_id}" - patched = client.patch(right, json={"purpose": "infra"}) + patched = await client.patch(right, json={"purpose": "infra"}) assert patched.status_code == 200 assert patched.json()["purpose"] == "infra" - assert client.delete(right).status_code == 204 - assert ProjectRepo.get(repo_id) is None + assert (await client.delete(right)).status_code == 204 + assert await ProjectRepo.get(repo_id) is None -def _make_work_item(project_id: int, ticket_key: str, *, resolved: bool = False): +async def _make_work_item(project_id: int, ticket_key: str, *, resolved: bool = False): from datetime import datetime from druks.contrib.ship.models import WorkItem - item = WorkItem.create( + item = await WorkItem.create( project_id=project_id, title=ticket_key, ticket_key=ticket_key, @@ -124,49 +136,51 @@ def _make_work_item(project_id: int, ticket_key: str, *, resolved: bool = False) return item -def test_deleting_a_project_cascades_its_work_items_and_spares_others(client: TestClient, druks_db): +async def test_deleting_a_project_cascades_its_work_items_and_spares_others( + client: TestClient, druks_db +): """DELETE cascades: the project and every work item it owns go, with no 409 reference guard — while another project's graph is left fully intact.""" from druks.contrib.ship.models import Project, WorkItem - target = Project.create(name="Target") - control = Project.create(name="Control") - doomed = _make_work_item(target.id, "ENG-1") - doomed_resolved = _make_work_item(target.id, "ENG-2", resolved=True) - survivor = _make_work_item(control.id, "ENG-3") + target = await Project.create(name="Target") + control = await Project.create(name="Control") + doomed = await _make_work_item(target.id, "ENG-1") + doomed_resolved = await _make_work_item(target.id, "ENG-2", resolved=True) + survivor = await _make_work_item(control.id, "ENG-3") target_id, control_id = target.id, control.id doomed_id, doomed_resolved_id, survivor_id = doomed.id, doomed_resolved.id, survivor.id - response = client.delete(f"/api/ship/projects/{target_id}") + response = await client.delete(f"/api/ship/projects/{target_id}") assert response.status_code == 204 - # The route committed on its own session; drop this session's identity map so - # the reads below reflect the committed graph rather than cached instances. - druks_db.expire_all() - assert Project.get(target_id) is None - assert WorkItem.get(doomed_id) is None - assert WorkItem.get(doomed_resolved_id) is None + # The route committed on its own session; drop the ambient session's identity + # map so the reads below reflect the committed graph, not cached instances. + db_session().expunge_all() + assert await Project.get(target_id) is None + assert await WorkItem.get(doomed_id) is None + assert await WorkItem.get(doomed_resolved_id) is None # The control project and its work item are untouched. - assert Project.get(control_id) is not None - assert WorkItem.get(survivor_id) is not None + assert await Project.get(control_id) is not None + assert await WorkItem.get(survivor_id) is not None -def test_the_repo_subject_read_side_mounts(client: TestClient, druks_db): +async def test_the_repo_subject_read_side_mounts(client: TestClient, druks_db): """Profile is about a repo, so the repo gets the board and the page it never had — its own runs' status and timeline, keyed by the repo's id.""" from druks.contrib.ship.models import Project, ProjectRepo from druks.contrib.ship.workflows import Profile from druks.testing import seed_run - project = Project.create(name="Acme") - repo = ProjectRepo.create(project_id=project.id, full_name="acme/widget") - seed_run(druks_db, kind=Profile.kind, subject=repo, state="running") + project = await Project.create(name="Acme") + repo = await ProjectRepo.create(project_id=project.id, full_name="acme/widget") + await seed_run(druks_db, kind=Profile.kind, subject=repo, state="running") - (row,) = client.get("/api/ship/project_repo").json()["rows"] + (row,) = (await client.get("/api/ship/project_repo")).json()["rows"] assert row["summary"]["id"] == str(repo.id) assert row["summary"]["fullName"] == "acme/widget" assert row["status"]["state"] == "running" - detail = client.get(f"/api/ship/project_repo/{repo.id}").json() + detail = (await client.get(f"/api/ship/project_repo/{repo.id}")).json() assert detail["summary"]["fullName"] == "acme/widget" assert [entry["kind"] for entry in detail["timeline"]] == [Profile.kind] diff --git a/backend/tests/ship/test_repo_routing.py b/backend/tests/ship/test_repo_routing.py index d3cdc2d1..38cee434 100644 --- a/backend/tests/ship/test_repo_routing.py +++ b/backend/tests/ship/test_repo_routing.py @@ -1,52 +1,52 @@ from druks.contrib.ship.models import Project, ProjectRepo -def _register(druks_db, *full_names): +async def _register(druks_db, *full_names): for full_name in full_names: - project = Project.create(name=full_name) - ProjectRepo.create(project_id=project.id, full_name=full_name) - druks_db.flush() + project = await Project.create(name=full_name) + await ProjectRepo.create(project_id=project.id, full_name=full_name) + await druks_db.flush() -def _lookup(**signals): +async def _lookup(**signals): defaults = {"project_name": None, "labels": []} - return ProjectRepo.lookup(**{**defaults, **signals}) + return await ProjectRepo.lookup(**{**defaults, **signals}) -def test_project_name_wins_over_labels(druks_db): - _register(druks_db, "acme/widget", "octo/alfred") - row = _lookup(project_name="widget", labels=["alfred"]) +async def test_project_name_wins_over_labels(druks_db): + await _register(druks_db, "acme/widget", "octo/alfred") + row = await _lookup(project_name="widget", labels=["alfred"]) assert row.full_name == "acme/widget" -def test_label_routes_when_project_name_is_not_a_repo(druks_db): +async def test_label_routes_when_project_name_is_not_a_repo(druks_db): """The org-project shape: the Jira project names the org, not a repo, and SHRP tickets carry a free-form 'Alfred' label — matched case-insensitively.""" - _register(druks_db, "octo/alfred") - row = _lookup(project_name="Octo", labels=["customer-request", "Alfred"]) + await _register(druks_db, "octo/alfred") + row = await _lookup(project_name="Octo", labels=["customer-request", "Alfred"]) assert row.full_name == "octo/alfred" -def test_first_matching_label_wins(druks_db): - _register(druks_db, "octo/alfred", "octo/obrv2") - row = _lookup(labels=["obrv2", "Alfred"]) +async def test_first_matching_label_wins(druks_db): + await _register(druks_db, "octo/alfred", "octo/obrv2") + row = await _lookup(labels=["obrv2", "Alfred"]) assert row.full_name == "octo/obrv2" -def test_no_signal_matches_any_repo(druks_db): - _register(druks_db, "octo/alfred") - assert _lookup(project_name="Octo", labels=["bug"]) is None +async def test_no_signal_matches_any_repo(druks_db): + await _register(druks_db, "octo/alfred") + assert await _lookup(project_name="Octo", labels=["bug"]) is None -def test_siblings_returns_only_other_repos_in_the_project(druks_db): - project = Project.create(name="Acme") - target = ProjectRepo.create(project_id=project.id, full_name="acme/api") - sibling = ProjectRepo.create( +async def test_siblings_returns_only_other_repos_in_the_project(druks_db): + project = await Project.create(name="Acme") + target = await ProjectRepo.create(project_id=project.id, full_name="acme/api") + sibling = await ProjectRepo.create( project_id=project.id, full_name="acme/web", purpose="frontend", ) - other_project = Project.create(name="Other") - ProjectRepo.create(project_id=other_project.id, full_name="other/worker") + other_project = await Project.create(name="Other") + await ProjectRepo.create(project_id=other_project.id, full_name="other/worker") - assert target.siblings() == [sibling] + assert await target.siblings() == [sibling] diff --git a/backend/tests/ship/test_subject_lifecycle.py b/backend/tests/ship/test_subject_lifecycle.py index b9355bbf..cab8edb7 100644 --- a/backend/tests/ship/test_subject_lifecycle.py +++ b/backend/tests/ship/test_subject_lifecycle.py @@ -15,11 +15,11 @@ pytestmark = pytest.mark.asyncio -def _work_item(**fields): - return make_test_work_item(repo="ClawHaven/acme-app", title="probe", **fields) +async def _work_item(**fields): + return await make_test_work_item(repo="ClawHaven/acme-app", title="probe", **fields) -def _subject_run( +async def _subject_run( druks_db, *, subject: WorkItem, @@ -35,8 +35,8 @@ def _subject_run( created_at=Base.utc_now() + timedelta(seconds=order), ) druks_db.add(run) - druks_db.flush() - seed_dbos_status(druks_db, run.id, state, subject=subject.identity) + await druks_db.flush() + await seed_dbos_status(druks_db, run.id, state, subject=subject.identity) return run @@ -44,15 +44,15 @@ async def test_gate_answer_resumes_only_a_run_parked_on_its_gate(druks_db, monke # A subject can carry runs of several workflows at once; the gate names which one # answers, so a newer run of another kind never hides the parked one. A timed-out # run keeps its stale ``input_gate``, so parked-ness decides, not that column. - subject = _work_item(ticket_key="ENG-748-A") - parked = _subject_run( + subject = await _work_item(ticket_key="ENG-748-A") + parked = await _subject_run( druks_db, subject=subject, kind=Build.kind, state="parked", gate=OperatorReply.name, ) - _subject_run(druks_db, subject=subject, kind=Profile.kind, state="running", order=1) + await _subject_run(druks_db, subject=subject, kind=Profile.kind, state="running", order=1) resumed = [] async def resume(self, **reply): @@ -63,8 +63,8 @@ async def resume(self, **reply): await OperatorReply.answer(subject, action="approve") assert resumed == [parked.id] - timed_out = _work_item(ticket_key="ENG-748-B") - _subject_run( + timed_out = await _work_item(ticket_key="ENG-748-B") + await _subject_run( druks_db, subject=timed_out, kind=Build.kind, @@ -81,9 +81,9 @@ async def test_workflow_cancel_takes_its_own_kind_and_passes_over_idle_subjects( ): # Webhooks redeliver, and a PR can close long after its build ended: cancelling what # is already gone is the no-op the caller expects, not an error. - subject = _work_item(ticket_key="ENG-748-C") - build = _subject_run(druks_db, subject=subject, kind=Build.kind, state="running") - _subject_run(druks_db, subject=subject, kind=Profile.kind, state="running", order=1) + subject = await _work_item(ticket_key="ENG-748-C") + build = await _subject_run(druks_db, subject=subject, kind=Build.kind, state="running") + await _subject_run(druks_db, subject=subject, kind=Profile.kind, state="running", order=1) cancelled = [] async def cancel(self, *, failure=None): @@ -94,8 +94,8 @@ async def cancel(self, *, failure=None): await Build.cancel(subject) assert cancelled == [build.id] - idle = _work_item(ticket_key="ENG-748-D") - _subject_run(druks_db, subject=idle, kind=Build.kind, state="finished") + idle = await _work_item(ticket_key="ENG-748-D") + await _subject_run(druks_db, subject=idle, kind=Build.kind, state="finished") await Build.cancel(idle) assert cancelled == [build.id] @@ -103,8 +103,8 @@ async def cancel(self, *, failure=None): async def test_cancel_and_answer_hold_the_caller_to_the_declared_subject(druks_db): # Build is about a work item; a repo names another timeline entirely, and a # non-subject names none. Both fail at the door rather than quietly no-opping. - item = _work_item(ticket_key="ENG-748-F") - repo = ProjectRepo.create(project_id=item.project_id, full_name="acme/app") + item = await _work_item(ticket_key="ENG-748-F") + repo = await ProjectRepo.create(project_id=item.project_id, full_name="acme/app") with pytest.raises(WorkflowError, match="is about WorkItem, not ProjectRepo"): await Build.cancel(repo) @@ -113,15 +113,17 @@ async def test_cancel_and_answer_hold_the_caller_to_the_declared_subject(druks_d async def test_subject_phase_reads_the_driving_running_workflow(druks_db, monkeypatch): - subject = _work_item(ticket_key="ENG-748-E") - _subject_run( + subject = await _work_item(ticket_key="ENG-748-E") + await _subject_run( druks_db, subject=subject, kind=Build.kind, state="parked", gate=OperatorReply.name, ) - driving = _subject_run(druks_db, subject=subject, kind=Profile.kind, state="running", order=1) + driving = await _subject_run( + druks_db, subject=subject, kind=Profile.kind, state="running", order=1 + ) seen = [] async def phase(workflow_id): diff --git a/backend/tests/ship/test_ticketing.py b/backend/tests/ship/test_ticketing.py index 9efdc4ed..b24f9478 100644 --- a/backend/tests/ship/test_ticketing.py +++ b/backend/tests/ship/test_ticketing.py @@ -18,19 +18,23 @@ def _pin_ship_settings(monkeypatch, **values): settings = Ship.Settings(**values) - monkeypatch.setattr(Ship, "settings", classmethod(lambda cls: settings)) + async def _settings(cls): + return settings -def _connect_linear(): - return ServiceIdentity.connect( + monkeypatch.setattr(Ship, "settings", classmethod(_settings)) + + +async def _connect_linear(): + return await ServiceIdentity.connect( "linear", identity={"actor": "druks", "workspace": "Acme"}, secrets={"api_key": "lin_secret", "webhook_secret": "lin-hook"}, ) -def _connect_jira(): - return ServiceIdentity.connect( +async def _connect_jira(): + return await ServiceIdentity.connect( "jira", identity={"base_url": "https://jira.test", "email": "a@b.com", "display_name": "druks"}, secrets={"api_token": "jira_secret", "webhook_secret": "jira-hook"}, @@ -40,15 +44,15 @@ def _connect_jira(): # --- Ship.get_tracker: the selected tracker ---------------------------------- -def test_tracker_builds_linear_from_the_service_row(druks_db, monkeypatch): - _connect_linear() +async def test_tracker_builds_linear_from_the_service_row(druks_db, monkeypatch): + await _connect_linear() _pin_ship_settings( monkeypatch, linear_resting_status="Backlog", linear_trigger_status="To Agent", ) - tracker = Ship.get_tracker("linear") + tracker = await Ship.get_tracker("linear") assert isinstance(tracker, Linear) assert tracker._client.api_key == "lin_secret" @@ -56,8 +60,8 @@ def test_tracker_builds_linear_from_the_service_row(druks_db, monkeypatch): assert tracker._status_names[TicketStatus.TRIGGER] == "To Agent" -def test_tracker_builds_jira_from_the_service_row(druks_db, monkeypatch): - _connect_jira() +async def test_tracker_builds_jira_from_the_service_row(druks_db, monkeypatch): + await _connect_jira() _pin_ship_settings( monkeypatch, tracker="jira", @@ -65,7 +69,7 @@ def test_tracker_builds_jira_from_the_service_row(druks_db, monkeypatch): jira_trigger_status="To Agent", ) - tracker = Ship.get_tracker("jira") + tracker = await Ship.get_tracker("jira") assert isinstance(tracker, Jira) assert tracker._client.base_url == "https://jira.test" @@ -73,30 +77,30 @@ def test_tracker_builds_jira_from_the_service_row(druks_db, monkeypatch): assert tracker._status_names[TicketStatus.TRIGGER] == "To Agent" -def test_tracker_is_none_for_github_and_a_disconnected_identity(druks_db, monkeypatch): - _connect_linear() +async def test_tracker_is_none_for_github_and_a_disconnected_identity(druks_db, monkeypatch): + await _connect_linear() _pin_ship_settings(monkeypatch) - assert not Ship.get_tracker("github") - assert not Ship.get_tracker("jira") + assert not await Ship.get_tracker("github") + assert not await Ship.get_tracker("jira") _pin_ship_settings(monkeypatch, tracker="jira") - assert not Ship.get_tracker("jira") - assert not Ship.get_tracker("linear") + assert not await Ship.get_tracker("jira") + assert not await Ship.get_tracker("linear") -def test_tracker_ignores_a_nonchosen_source_with_a_connected_identity(druks_db, monkeypatch): - _connect_linear() +async def test_tracker_ignores_a_nonchosen_source_with_a_connected_identity(druks_db, monkeypatch): + await _connect_linear() _pin_ship_settings(monkeypatch, tracker="jira") - assert not Ship.get_tracker("linear") + assert not await Ship.get_tracker("linear") -def test_empty_resting_status_leaves_backlog_unmapped(druks_db, monkeypatch): - _connect_linear() +async def test_empty_resting_status_leaves_backlog_unmapped(druks_db, monkeypatch): + await _connect_linear() _pin_ship_settings(monkeypatch, linear_resting_status="") - tracker = Ship.get_tracker("linear") + tracker = await Ship.get_tracker("linear") assert TicketStatus.BACKLOG not in tracker._status_names @@ -174,29 +178,29 @@ async def fake_get(self, url, **kwargs): # --- The tracker doctor check ------------------------------------------------- -def test_tracker_check_accepts_trackerless_by_choice(monkeypatch): +async def test_tracker_check_accepts_trackerless_by_choice(monkeypatch): _pin_ship_settings(monkeypatch, tracker="none") - result = check_tracker_identity() + result = await check_tracker_identity() assert result.ok assert "choice" in result.detail -def test_tracker_check_reports_a_selected_connected_tracker(druks_db, monkeypatch): - _connect_linear() +async def test_tracker_check_reports_a_selected_connected_tracker(druks_db, monkeypatch): + await _connect_linear() _pin_ship_settings(monkeypatch) - result = check_tracker_identity() + result = await check_tracker_identity() assert result.ok assert "linear" in result.detail -def test_tracker_check_pends_a_selected_unconnected_tracker(druks_db, monkeypatch): +async def test_tracker_check_pends_a_selected_unconnected_tracker(druks_db, monkeypatch): _pin_ship_settings(monkeypatch, tracker="jira") - result = check_tracker_identity() + result = await check_tracker_identity() assert not result.ok assert result.pending @@ -345,6 +349,13 @@ def handler(request: httpx.Request) -> httpx.Response: # --- WorkItem.set_ticket_status: the status-push consumer ------------------- +def _tracker_stub(fake): + async def get_tracker(cls, source=None): + return fake + + return classmethod(get_tracker) + + class _FakeTracker: known_exceptions: tuple = () @@ -366,9 +377,11 @@ async def aclose(self): @pytest.mark.asyncio async def test_ticket_state_pushes_status(druks_db, monkeypatch): - item = make_test_work_item(repo="acme/widget", source="linear", ticket_key="ACME-1", title="t") + item = await make_test_work_item( + repo="acme/widget", source="linear", ticket_key="ACME-1", title="t" + ) fake = _FakeTracker() - monkeypatch.setattr(Ship, "get_tracker", classmethod(lambda cls, source=None: fake)) + monkeypatch.setattr(Ship, "get_tracker", _tracker_stub(fake)) await item.set_ticket_status(TicketStatus.DONE) @@ -377,14 +390,18 @@ async def test_ticket_state_pushes_status(druks_db, monkeypatch): @pytest.mark.asyncio async def test_ticket_state_skips_non_tracker_source(druks_db): - item = make_test_work_item(repo="acme/widget", source="github", ticket_key="#5", title="t") + item = await make_test_work_item( + repo="acme/widget", source="github", ticket_key="#5", title="t" + ) # github has no tracker — a no-op that must not raise. await item.set_ticket_status(TicketStatus.DONE) @pytest.mark.asyncio async def test_ticket_state_closes_on_failure(druks_db, monkeypatch): - item = make_test_work_item(repo="acme/widget", source="linear", ticket_key="ACME-2", title="t") + item = await make_test_work_item( + repo="acme/widget", source="linear", ticket_key="ACME-2", title="t" + ) class _Boom(_FakeTracker): known_exceptions = (LinearAPIError,) @@ -393,7 +410,7 @@ async def set_status(self, key, status): raise LinearAPIError("boom") boom = _Boom() - monkeypatch.setattr(Ship, "get_tracker", classmethod(lambda cls, source=None: boom)) + monkeypatch.setattr(Ship, "get_tracker", _tracker_stub(boom)) await item.set_ticket_status(TicketStatus.DONE) diff --git a/backend/tests/ship/test_webhooks_jira.py b/backend/tests/ship/test_webhooks_jira.py index adbbac76..d97e68be 100644 --- a/backend/tests/ship/test_webhooks_jira.py +++ b/backend/tests/ship/test_webhooks_jira.py @@ -58,40 +58,40 @@ def test_route_is_unchanged(): assert f"{webhooks_router.prefix}/{JiraEvents.path}" == "/_external/jira/events/" -def _connect_jira(*, base_url="https://jira.test/", webhook_secret="s3cret"): - return ServiceIdentity.connect( +async def _connect_jira(*, base_url="https://jira.test/", webhook_secret="s3cret"): + return await ServiceIdentity.connect( "jira", identity={"base_url": base_url, "email": "a@b.com", "display_name": "druks"}, secrets={"api_token": "tok", "webhook_secret": webhook_secret}, ) -def test_rejects_when_not_connected(tmp_path, druks_db): +async def test_rejects_when_not_connected(tmp_path, druks_db): events = _provider(tmp_path, payload=_issue()) with pytest.raises(HTTPException) as exc: - events.request_is_authentic() + await events.request_is_authentic() assert exc.value.status_code == 401 -def test_authentic_when_token_header_matches(tmp_path, druks_db): - _connect_jira() +async def test_authentic_when_token_header_matches(tmp_path, druks_db): + await _connect_jira() events = _provider( tmp_path, payload=_issue(), headers={"x-druks-webhook-token": "s3cret"}, ) - assert events.request_is_authentic() + assert await events.request_is_authentic() -def test_rejects_when_token_missing_or_wrong(tmp_path, druks_db): - _connect_jira() +async def test_rejects_when_token_missing_or_wrong(tmp_path, druks_db): + await _connect_jira() events = _provider( tmp_path, payload=_issue(), headers={"x-druks-webhook-token": "nope"}, ) with pytest.raises(HTTPException) as exc: - events.request_is_authentic() + await events.request_is_authentic() assert exc.value.status_code == 401 @@ -107,7 +107,7 @@ async def test_emits_normalized_ticket_transition(tmp_path, druks_db, monkeypatc async def _emit(event_type, **kwargs): captured.update({"event": event_type, **kwargs}) - _connect_jira() + await _connect_jira() monkeypatch.setattr(webhook_module, "publish", _emit) await _provider(tmp_path, payload=_issue(key="IT-9", status="Ready")).on_issue_event() @@ -127,7 +127,7 @@ async def test_done_category_marks_the_transition_terminal(tmp_path, druks_db, m async def _emit(event_type, **kwargs): events.append((event_type, kwargs["payload"])) - _connect_jira() + await _connect_jira() monkeypatch.setattr(webhook_module, "publish", _emit) payload = _issue(key="IT-9", status="Done", status_category="done") await _provider(tmp_path, payload=payload).on_issue_event() @@ -143,7 +143,7 @@ async def test_open_category_is_not_terminal(tmp_path, druks_db, monkeypatch): async def _emit(event_type, **kwargs): events.append((event_type, kwargs["payload"])) - _connect_jira() + await _connect_jira() monkeypatch.setattr(webhook_module, "publish", _emit) provider = _provider( tmp_path, payload=_issue(status="In Progress", status_category="indeterminate") @@ -159,7 +159,11 @@ async def _emit(event_type, **kwargs): def _pin_settings(monkeypatch, **over): settings = subs.Ship.Settings(**{"tracker": "jira", **over}) - monkeypatch.setattr(subs.Ship, "settings", classmethod(lambda cls: settings)) + + async def _settings(cls): + return settings + + monkeypatch.setattr(subs.Ship, "settings", classmethod(_settings)) async def test_trigger_status_dispatches_build_with_the_webhook_payload(tmp_path, monkeypatch): @@ -175,14 +179,14 @@ async def test_trigger_status_dispatches_build_with_the_webhook_payload(tmp_path async def test_trigger_status_does_not_redispatch_a_merged_item(druks_db, monkeypatch): - item = make_test_work_item( + item = await make_test_work_item( repo="octo/alfred", source="jira", ticket_key="IT-12", title="Add an endpoint", ) item.resolution = "merged" - druks_db.flush() + await druks_db.flush() _pin_settings(monkeypatch, jira_trigger_status="Ready") start = AsyncMock() monkeypatch.setattr(subs.Build, "start", start) @@ -193,19 +197,19 @@ async def test_trigger_status_does_not_redispatch_a_merged_item(druks_db, monkey async def test_trigger_status_redispatches_a_closed_item(druks_db, monkeypatch): - ServiceIdentity.connect( + await ServiceIdentity.connect( "github", identity={"app_id": "1", "slug": "druks-operator"}, secrets={"private_key": "operator-pem", "webhook_secret": "hook-secret"}, ) - item = make_test_work_item( + item = await make_test_work_item( repo="octo/alfred", source="jira", ticket_key="IT-12", title="Add an endpoint", ) item.resolution = "closed" - druks_db.flush() + await druks_db.flush() _pin_settings(monkeypatch, jira_trigger_status="Ready") start = AsyncMock() monkeypatch.setattr(subs.Build, "start", start) @@ -219,11 +223,11 @@ async def test_trigger_status_routes_a_new_ticket_by_label(tmp_path, druks_db, m """No work item yet: the label names the repo, the registry routes it.""" from druks.contrib.ship.models import Project, ProjectRepo, WorkItem - project = Project.create(name="octo/alfred") - ProjectRepo.create(project_id=project.id, full_name="octo/alfred") - druks_db.flush() + project = await Project.create(name="octo/alfred") + await ProjectRepo.create(project_id=project.id, full_name="octo/alfred") + await druks_db.flush() _pin_settings(monkeypatch, jira_trigger_status="Ready") - seed_run(druks_db, kind=Build.kind, run_id="run-new") + await seed_run(druks_db, kind=Build.kind, run_id="run-new") async def fake_start(cls, **kwargs): return "run-new" @@ -234,7 +238,7 @@ async def fake_start(cls, **kwargs): payload=_jira_payload(key="SHRP-1", status="Ready", project="Octo", labels=["Alfred"]), ) - item = WorkItem.get_for_ticket_key(source="jira", ticket_key="SHRP-1") + item = await WorkItem.get_for_ticket_key(source="jira", ticket_key="SHRP-1") assert item.repo == "octo/alfred" assert item.project_id == project.id diff --git a/backend/tests/ship/test_webhooks_linear.py b/backend/tests/ship/test_webhooks_linear.py index d95edac0..db0cb059 100644 --- a/backend/tests/ship/test_webhooks_linear.py +++ b/backend/tests/ship/test_webhooks_linear.py @@ -52,11 +52,11 @@ def test_route_is_unchanged(): assert f"{webhooks_router.prefix}/{LinearEvents.path}" == "/_external/linear/events/" -def test_authentication_reads_the_service_row(tmp_path, druks_db): +async def test_authentication_reads_the_service_row(tmp_path, druks_db): secret = "linear-secret" raw_body = b"{}" signature = hmac.new(secret.encode(), raw_body, hashlib.sha256).hexdigest() - ServiceIdentity.connect( + await ServiceIdentity.connect( "linear", identity={"actor": "druks", "workspace": "Acme"}, secrets={"api_key": "lin_secret", "webhook_secret": secret}, @@ -68,10 +68,10 @@ def test_authentication_reads_the_service_row(tmp_path, druks_db): ) events.raw_body = raw_body - assert events.request_is_authentic() + assert await events.request_is_authentic() -def test_rejects_when_not_connected(tmp_path, druks_db): +async def test_rejects_when_not_connected(tmp_path, druks_db): events = _provider( tmp_path, payload={}, @@ -80,7 +80,7 @@ def test_rejects_when_not_connected(tmp_path, druks_db): events.raw_body = b"{}" with pytest.raises(HTTPException) as error: - events.request_is_authentic() + await events.request_is_authentic() assert error.value.status_code == 401 assert "not connected" in error.value.detail diff --git a/backend/tests/ship/test_webhooks_pull_request.py b/backend/tests/ship/test_webhooks_pull_request.py index 7fa34d13..ebd46b0f 100644 --- a/backend/tests/ship/test_webhooks_pull_request.py +++ b/backend/tests/ship/test_webhooks_pull_request.py @@ -26,10 +26,10 @@ async def _fetch(*, repo, path): monkeypatch.setattr("druks.apps.config.fetch_file", _fetch) -def _milestone_count(work_item_id, milestone): +async def _milestone_count(work_item_id, milestone): from druks.database import db_session - return db_session().scalar( + return await db_session().scalar( select(func.count()) .select_from(Event) .where( @@ -60,14 +60,14 @@ async def _fire_closed(*, repo, pr_number, branch, tmp_path, merged=True, at=_RE await events.on_pull_request_closed() -def _park_work_item(*, repo, pr_number, branch, state="parked", input_gate="review_work"): +async def _park_work_item(*, repo, pr_number, branch, state="parked", input_gate="review_work"): """A work item with a build run paused on the operator (review_work) — the haunting case. Returns (work_item_id, run_id).""" from druks.database import db_session - item = make_test_work_item(repo=repo, title="Externally merged") - item.update(pr_number=pr_number, branch=branch) - run = seed_build_run( + item = await make_test_work_item(repo=repo, title="Externally merged") + await item.update(pr_number=pr_number, branch=branch) + run = await seed_build_run( db_session(), work_item_id=item.id, state=state, @@ -76,29 +76,29 @@ def _park_work_item(*, repo, pr_number, branch, state="parked", input_gate="revi return item.id, run.id -def _fresh_run(run_id): +async def _fresh_run(run_id): # Workflow.cancel() never writes state — re-select before reading the derived one. from druks.database import db_session - db_session().expire_all() - return Run.get(run_id) + db_session().expunge_all() + return await Run.get(run_id) @pytest.mark.asyncio async def test_external_merge_stores_githubs_verdict_and_ends_involvement(druks_db, tmp_path): repo, pr_number, branch = "ClawHaven/acme-app", 42, "agent/eng-1" - work_item_id, run_id = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, run_id = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) # The verdict is stored with GitHub's own stamp, not druks's receipt time... - item = WorkItem.get(work_item_id) + item = await WorkItem.get(work_item_id) assert item.resolution == "merged" assert item.resolved_at == datetime(2026, 7, 25, 21, 59, 9, tzinfo=UTC) # ...announced as the milestone... - assert _milestone_count(work_item_id, "merged") == 1 + assert await _milestone_count(work_item_id, "merged") == 1 # ...and involvement ended: the parked build run is cancelled. - assert not _fresh_run(run_id).is_active + assert not (await _fresh_run(run_id)).is_active @pytest.mark.asyncio @@ -107,7 +107,7 @@ async def test_merge_ships_but_leaves_a_running_build_to_converge(druks_db, tmp_ druks's own merges too. A RUNNING run is left alone: it converges on its own (its merge step sees the closed PR).""" repo, pr_number, branch = "ClawHaven/acme-app", 43, "agent/eng-2" - work_item_id, run_id = _park_work_item( + work_item_id, run_id = await _park_work_item( repo=repo, pr_number=pr_number, branch=branch, @@ -116,8 +116,8 @@ async def test_merge_ships_but_leaves_a_running_build_to_converge(druks_db, tmp_ await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) - assert Run.get(run_id).state == "running" # not cancelled from under druks - assert _milestone_count(work_item_id, "merged") == 1 + assert (await Run.get(run_id)).state == "running" # not cancelled from under druks + assert await _milestone_count(work_item_id, "merged") == 1 @pytest.mark.asyncio @@ -126,7 +126,7 @@ async def test_a_redelivered_webhook_does_not_rewrite_the_verdict(druks_db, tmp_ same path. The stored verdict is the first one; a contradicting redelivery neither overwrites it nor records a second milestone.""" repo, pr_number, branch = "ClawHaven/acme-app", 44, "agent/eng-3" - work_item_id, _ = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, _ = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) await _fire_closed( @@ -138,11 +138,11 @@ async def test_a_redelivered_webhook_does_not_rewrite_the_verdict(druks_db, tmp_ at="2026-07-26T13:30:30Z", ) - item = WorkItem.get(work_item_id) + item = await WorkItem.get(work_item_id) assert item.resolution == "merged" assert item.resolved_at == datetime(2026, 7, 25, 21, 59, 9, tzinfo=UTC) - assert _milestone_count(work_item_id, "merged") == 1 - assert _milestone_count(work_item_id, "closed") == 0 + assert await _milestone_count(work_item_id, "merged") == 1 + assert await _milestone_count(work_item_id, "closed") == 0 @pytest.mark.asyncio @@ -151,7 +151,7 @@ async def test_closed_unmerged_stores_closed_and_ends_involvement(druks_db, tmp_ the branch). Store GitHub's 'closed' and un-park, so the item leaves the active board for History rather than being ignored.""" repo, pr_number, branch = "ClawHaven/acme-app", 45, "agent/eng-4" - work_item_id, run_id = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, run_id = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, @@ -161,16 +161,16 @@ async def test_closed_unmerged_stores_closed_and_ends_involvement(druks_db, tmp_ merged=False, ) - assert WorkItem.get(work_item_id).resolution == "closed" - assert _milestone_count(work_item_id, "closed") == 1 - assert _milestone_count(work_item_id, "merged") == 0 - assert not _fresh_run(run_id).is_active + assert (await WorkItem.get(work_item_id)).resolution == "closed" + assert await _milestone_count(work_item_id, "closed") == 1 + assert await _milestone_count(work_item_id, "merged") == 0 + assert not (await _fresh_run(run_id)).is_active @pytest.mark.asyncio async def test_closed_unmerged_cancels_in_flight_run(druks_db, tmp_path): repo, pr_number, branch = "ClawHaven/acme-app", 46, "agent/eng-5" - work_item_id, run_id = _park_work_item( + work_item_id, run_id = await _park_work_item( repo=repo, pr_number=pr_number, branch=branch, @@ -181,8 +181,8 @@ async def test_closed_unmerged_cancels_in_flight_run(druks_db, tmp_path): repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=False ) - assert _fresh_run(run_id).state == "cancelled" - assert _milestone_count(work_item_id, "closed") == 1 + assert (await _fresh_run(run_id)).state == "cancelled" + assert await _milestone_count(work_item_id, "closed") == 1 @pytest.mark.asyncio @@ -191,7 +191,7 @@ async def test_a_merge_after_a_failed_build_still_settles_the_item(druks_db, tmp nothing in the run lifecycle can announce the operator's later manual merge. GitHub's does, and the stored verdict takes the item to History.""" repo, pr_number, branch = "ClawHaven/acme-app", 126, "agent/eng-760" - work_item_id, _ = _park_work_item( + work_item_id, _ = await _park_work_item( repo=repo, pr_number=pr_number, branch=branch, @@ -200,10 +200,10 @@ async def test_a_merge_after_a_failed_build_still_settles_the_item(druks_db, tmp await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) - item = WorkItem.get(work_item_id) + item = await WorkItem.get(work_item_id) assert item.resolution == "merged" - assert item.id not in {summary.id for summary in WorkItem.list_summaries(None)} - assert [row.id for row in WorkItem.list_handoff()] == [item.id] + assert item.id not in {summary.id for summary in await WorkItem.list_summaries(None)} + assert [row.id for row in await WorkItem.list_handoff()] == [item.id] @pytest.mark.asyncio @@ -214,17 +214,17 @@ async def test_a_remerge_after_redispatch_records_a_fresh_verdict(druks_db, tmp_ from druks.database import db_session as ds repo, pr_number, branch = "ClawHaven/acme-app", 77, "agent/eng-9" - work_item_id, _ = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, _ = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) # Redispatched: a newer build run owns the item, and its PR is a fresh one. - seed_build_run(ds(), work_item_id=work_item_id, state="running") - WorkItem.get(work_item_id).start_attempt() - WorkItem.get(work_item_id).update(pr_number=pr_number, branch=branch) + await seed_build_run(ds(), work_item_id=work_item_id, state="running") + await (await WorkItem.get(work_item_id)).start_attempt() + await (await WorkItem.get(work_item_id)).update(pr_number=pr_number, branch=branch) await _fire_closed(repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path) - assert WorkItem.get(work_item_id).resolution == "merged" - assert _milestone_count(work_item_id, "merged") == 2 + assert (await WorkItem.get(work_item_id)).resolution == "merged" + assert await _milestone_count(work_item_id, "merged") == 2 @pytest.mark.asyncio @@ -243,7 +243,7 @@ async def _record(self, status): monkeypatch.setattr(WorkItem, "set_ticket_status", _record) repo, pr_number, branch = "ClawHaven/acme-app", 91, "agent/eng-20" - work_item_id, _ = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, _ = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=False @@ -266,7 +266,7 @@ async def _record(self, status): monkeypatch.setattr(WorkItem, "set_ticket_status", _record) repo, pr_number, branch = "ClawHaven/acme-app", 92, "agent/eng-21" - work_item_id, _ = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, _ = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=True @@ -292,11 +292,11 @@ async def _record(repo, branch): deleted.append((repo, branch)) monkeypatch.setattr( - build_models, "get_github_client", lambda: SimpleNamespace(delete_branch=_record) + build_models, "get_github_client", _async_value(SimpleNamespace(delete_branch=_record)) ) repo, pr_number, branch = "ClawHaven/acme-app", 93, "agent/eng-22" - _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=False @@ -305,6 +305,13 @@ async def _record(repo, branch): assert deleted == [] +def _async_value(value): + async def get(): + return value + + return get + + @pytest.mark.asyncio async def test_external_close_deletes_branch_by_default(druks_db, tmp_path, monkeypatch): from druks.contrib.ship import models as build_models @@ -315,11 +322,11 @@ async def _record(repo, branch): deleted.append((repo, branch)) monkeypatch.setattr( - build_models, "get_github_client", lambda: SimpleNamespace(delete_branch=_record) + build_models, "get_github_client", _async_value(SimpleNamespace(delete_branch=_record)) ) repo, pr_number, branch = "ClawHaven/acme-app", 94, "agent/eng-23" - _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=False @@ -358,7 +365,7 @@ async def _record(self, status): monkeypatch.setattr(WorkItem, "set_ticket_status", _record) repo, pr_number, branch = "ClawHaven/acme-app", 95, "agent/eng-24" - work_item_id, _ = _park_work_item(repo=repo, pr_number=pr_number, branch=branch) + work_item_id, _ = await _park_work_item(repo=repo, pr_number=pr_number, branch=branch) await _fire_closed( repo=repo, pr_number=pr_number, branch=branch, tmp_path=tmp_path, merged=False @@ -366,7 +373,7 @@ async def _record(self, status): assert deleted == [] # cleanup skipped when policy can't be resolved assert pushed == [TicketStatus.BACKLOG] # ticket still reset - assert WorkItem.get(work_item_id).resolution == "closed" + assert (await WorkItem.get(work_item_id)).resolution == "closed" @pytest.mark.asyncio @@ -377,13 +384,13 @@ async def test_stale_close_after_redispatch_spares_the_new_run(druks_db, tmp_pat from druks.database import db_session as ds repo, pr_a, branch_a = "ClawHaven/acme-app", 61, "agent/eng-old" - item = make_test_work_item(repo=repo, title="Re-dispatched") - item.update(pr_number=pr_a, branch=branch_a) + item = await make_test_work_item(repo=repo, title="Re-dispatched") + await item.update(pr_number=pr_a, branch=branch_a) # Re-dispatch: a new run takes over and claims the item. - new_run = seed_build_run(ds(), work_item_id=item.id, state="running") - item.start_attempt() + new_run = await seed_build_run(ds(), work_item_id=item.id, state="running") + await item.start_attempt() await _fire_closed(repo=repo, pr_number=pr_a, branch=branch_a, tmp_path=tmp_path, merged=False) - assert _fresh_run(new_run.id).state == "running" # the live attempt is untouched + assert (await _fresh_run(new_run.id)).state == "running" # the live attempt is untouched assert item.resolution is None diff --git a/backend/tests/ship/test_webhooks_push.py b/backend/tests/ship/test_webhooks_push.py index ced3856c..981e7844 100644 --- a/backend/tests/ship/test_webhooks_push.py +++ b/backend/tests/ship/test_webhooks_push.py @@ -37,8 +37,8 @@ async def _dispatch(cls, repo, *, refresh_only=False): async def test_policy_push_on_default_branch_reprofiles(tmp_path, druks_db, monkeypatch): from druks.contrib.ship.models import Project, ProjectRepo - project = Project.create(name="Acme") - repo = ProjectRepo.create(project_id=project.id, full_name="acme/widget") + project = await Project.create(name="Acme") + repo = await ProjectRepo.create(project_id=project.id, full_name="acme/widget") calls = _stub_profile_dispatch(monkeypatch) await _fire_push( @@ -62,8 +62,8 @@ async def test_policy_push_on_default_branch_reprofiles(tmp_path, druks_db, monk async def test_non_default_branch_push_is_ignored(tmp_path, druks_db, monkeypatch): from druks.contrib.ship.models import Project, ProjectRepo - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name="acme/widget") + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name="acme/widget") calls = _stub_profile_dispatch(monkeypatch) await _fire_push( @@ -82,8 +82,8 @@ async def test_non_default_branch_push_is_ignored(tmp_path, druks_db, monkeypatc async def test_unrelated_path_push_is_ignored(tmp_path, druks_db, monkeypatch): from druks.contrib.ship.models import Project, ProjectRepo - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name="acme/widget") + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name="acme/widget") calls = _stub_profile_dispatch(monkeypatch) await _fire_push( diff --git a/backend/tests/test_agent_call_liveness.py b/backend/tests/test_agent_call_liveness.py index 04c7c218..edef22da 100644 --- a/backend/tests/test_agent_call_liveness.py +++ b/backend/tests/test_agent_call_liveness.py @@ -1,36 +1,38 @@ from druks.durable.dbos_state import workflow_status from druks.durable.enums import AgentCallStatus +from druks.durable.models import AgentCall from druks.testing import seed_call, seed_run from druks_field_notes.models import Note from druks_field_notes.workflows import Summarize from sqlalchemy import update -def _running_call(druks_db): - note = Note.create(body="agent call liveness") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - return seed_call(druks_db, run, "summarize", status="running") +async def _running_call(druks_db): + note = await Note.create(body="agent call liveness") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + return await seed_call(druks_db, run, "summarize", status="running") -def test_an_unfinished_call_reads_running_while_its_run_is_active(druks_db): - assert _running_call(druks_db).live_status == AgentCallStatus.RUNNING +async def test_an_unfinished_call_reads_running_while_its_run_is_active(druks_db): + call = await AgentCall.get((await _running_call(druks_db)).id) + assert call.live_status == AgentCallStatus.RUNNING -def test_an_unfinished_call_reads_abandoned_once_its_run_is_terminal(druks_db): - call = _running_call(druks_db) - druks_db.execute( +async def test_an_unfinished_call_reads_abandoned_once_its_run_is_terminal(druks_db): + call = await _running_call(druks_db) + await druks_db.execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == call.run_id) .values(status="ERROR") ) - druks_db.expire_all() + call = await AgentCall.get(call.id) assert call.live_status == AgentCallStatus.ABANDONED -def test_a_finished_call_keeps_its_outcome(druks_db): - note = Note.create(body="finished agent call") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call( +async def test_a_finished_call_keeps_its_outcome(druks_db): + note = await Note.create(body="finished agent call") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call( druks_db, run, "summarize", diff --git a/backend/tests/test_agent_routes.py b/backend/tests/test_agent_routes.py index e00ce7bc..6c661a2e 100644 --- a/backend/tests/test_agent_routes.py +++ b/backend/tests/test_agent_routes.py @@ -46,14 +46,14 @@ def client(tmp_path: Path, druks_db, monkeypatch): @pytest.fixture -def account(druks_db): +async def account(druks_db): # The account configure_app_for_test signs requests in as. - return Account.get_or_create("op@example.com") + return await Account.get_or_create("op@example.com") -def _connect_github() -> None: +async def _connect_github() -> None: # Start routes guard on the GitHub service identity before spending a run. - ServiceIdentity.connect( + await ServiceIdentity.connect( "github", identity={"app_id": "1", "slug": "druks-operator"}, secrets={"private_key": "operator-pem", "webhook_secret": "hook-secret"}, @@ -71,11 +71,11 @@ async def _spy(self, **fields): return calls -def _park(druks_db, note, *, context: str = ""): +async def _park(druks_db, note, *, context: str = ""): ask = dict(_IN_APP_ASK) if context: ask["context"] = context - run = seed_run( + run = await seed_run( druks_db, kind=Summarize.kind, subject=note, @@ -84,7 +84,7 @@ def _park(druks_db, note, *, context: str = ""): input_request=ask, ) run.input_requested_at = datetime.now(UTC) - druks_db.flush() + await druks_db.flush() return run @@ -124,12 +124,12 @@ def test_openapi_pins_platform_and_app_agent_routes(client: TestClient): } -def test_review_request_returns_the_run_id_start_hands_back( +async def test_review_request_returns_the_run_id_start_hands_back( client: TestClient, account: Account, monkeypatch ): - _connect_github() - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name="acme/app") + await _connect_github() + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name="acme/app") live_run_id = "review-run-id" starts = [] @@ -178,23 +178,30 @@ async def aclose(self): self.calls.append("aclose") -def test_ship_start_stamps_the_trigger_status_for_known_and_unknown_tickets( +def _tracker_stub(fake): + async def get_tracker(cls, source=None): + return fake + + return classmethod(get_tracker) + + +async def test_ship_start_stamps_the_trigger_status_for_known_and_unknown_tickets( client: TestClient, account: Account, monkeypatch ): # ENG-831 has a local work item, ENG-777 has never been seen — both take # the same tracker path; webhook intake, not the route, opens builds. - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name="acme/app") - WorkItem.create( + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name="acme/app") + await WorkItem.create( project_id=project.id, source="linear", title="Build the agent route", ticket_key="ENG-831", repo="acme/app", ) - _, pat_token = PersonalAccessToken.create(account_id=account.id, name="agent") + _, pat_token = await PersonalAccessToken.create(account_id=account.id, name="agent") fake = _FakeTracker() - monkeypatch.setattr(Ship, "get_tracker", classmethod(lambda cls, source=None: fake)) + monkeypatch.setattr(Ship, "get_tracker", _tracker_stub(fake)) starts = [] async def start(cls, **kwargs): @@ -226,7 +233,7 @@ async def start(cls, **kwargs): def test_ship_start_translates_an_unknown_tracker_ticket(client: TestClient, monkeypatch): fake = _FakeTracker(error=UnknownTicketError("ENG-9999", "Linear")) - monkeypatch.setattr(Ship, "get_tracker", classmethod(lambda cls, source=None: fake)) + monkeypatch.setattr(Ship, "get_tracker", _tracker_stub(fake)) response = client.post("/api/ship/work-items/ENG-9999/start") @@ -254,7 +261,7 @@ def test_ship_start_does_not_acknowledge_a_tracker_failure(tmp_path, druks_db, m monkeypatch.setenv("DRUKS_DATA_DIR", str(tmp_path)) app = configure_app_for_test(settings=make_settings(tmp_path)) fake = _FakeTracker(error=LinearAPIError("linear fell over")) - monkeypatch.setattr(Ship, "get_tracker", classmethod(lambda cls, source=None: fake)) + monkeypatch.setattr(Ship, "get_tracker", _tracker_stub(fake)) with TestClient(app, raise_server_exceptions=False) as failing: response = failing.post("/api/ship/work-items/ENG-831/start") @@ -263,9 +270,9 @@ def test_ship_start_does_not_acknowledge_a_tracker_failure(tmp_path, druks_db, m assert response.json() == {"error": "INTERNAL_ERROR", "detail": "Internal server error"} -def test_review_request_refuses_when_github_is_not_connected(client: TestClient, monkeypatch): - project = Project.create(name="Acme") - ProjectRepo.create(project_id=project.id, full_name="acme/app") +async def test_review_request_refuses_when_github_is_not_connected(client: TestClient, monkeypatch): + project = await Project.create(name="Acme") + await ProjectRepo.create(project_id=project.id, full_name="acme/app") async def start(cls, **kwargs): raise AssertionError("start must not be reached without a GitHub identity") @@ -294,7 +301,7 @@ def test_agent_routes_sit_behind_the_gate(tmp_path, druks_db): assert anonymous.post("/api/ship/work-items/ENG-831/start").status_code == 401 -def test_agent_errors_share_one_shape(client: TestClient, druks_db): +async def test_agent_errors_share_one_shape(client: TestClient, druks_db): missing = client.get("/api/gates/no-such-run") assert missing.status_code == 404 assert missing.json() == { @@ -303,8 +310,8 @@ def test_agent_errors_share_one_shape(client: TestClient, druks_db): "retryable": False, } - note = Note.create(body="stale gate") - run = _park(druks_db, note) + note = await Note.create(body="stale gate") + run = await _park(druks_db, note) stale = client.post( f"/api/gates/{run.id}/answer", json={"parkedAt": "2020-01-01T00:00:00+00:00", "control": "approve"}, @@ -326,16 +333,18 @@ def test_missing_agent_call_uses_the_unified_shape(client: TestClient, druks_db) } -def test_list_open_subjects_returns_newest_open_work_and_latest_calls(client: TestClient, druks_db): - finished_note = Note.create(body="finished") - seed_run(druks_db, kind=Summarize.kind, subject=finished_note, state="finished") +async def test_list_open_subjects_returns_newest_open_work_and_latest_calls( + client: TestClient, druks_db +): + finished_note = await Note.create(body="finished") + await seed_run(druks_db, kind=Summarize.kind, subject=finished_note, state="finished") - failed_note = Note.create(body="failed") - older = seed_run(druks_db, kind=Summarize.kind, subject=failed_note) + failed_note = await Note.create(body="failed") + older = await seed_run(druks_db, kind=Summarize.kind, subject=failed_note) older.created_at = datetime(2026, 1, 1, tzinfo=UTC) - seed_call(druks_db, older, "older") + await seed_call(druks_db, older, "older") failure = "discarded failure prefix " + "f" * 512 - newest = seed_run( + newest = await seed_run( druks_db, kind=Summarize.kind, subject=failed_note, @@ -343,10 +352,10 @@ def test_list_open_subjects_returns_newest_open_work_and_latest_calls(client: Te failure=failure, ) newest.created_at = older.created_at + timedelta(days=1) - seed_call(druks_db, newest, "first") - latest_call = seed_call(druks_db, newest, "latest") + await seed_call(druks_db, newest, "first") + latest_call = await seed_call(druks_db, newest, "latest") subject_label = "long label kept whole " + "l" * 512 - druks_db.execute( + await druks_db.execute( workflow_status.update() .where(workflow_status.c.workflow_uuid == newest.id) .values( @@ -358,9 +367,9 @@ def test_list_open_subjects_returns_newest_open_work_and_latest_calls(client: Te ) ) - callless_note = Note.create(body="callless") - seed_run(druks_db, kind="field_notes.audit", subject=callless_note) - seed_run(druks_db, kind="usage.scrape") + callless_note = await Note.create(body="callless") + await seed_run(druks_db, kind="field_notes.audit", subject=callless_note) + await seed_run(druks_db, kind="usage.scrape") response = client.get("/api/open-subjects") @@ -386,11 +395,11 @@ def test_list_open_subjects_returns_newest_open_work_and_latest_calls(client: Te assert subjects[str(callless_note.id)]["workflows"][0]["latestAgentCall"] is None -def test_list_open_subjects_keeps_type_and_kind_partitions(client: TestClient, druks_db): - typed_note = Note.create(body="two types") - note_run = seed_run(druks_db, kind=Summarize.kind, subject=typed_note) - ticket_run = seed_run(druks_db, kind=Summarize.kind, subject=typed_note) - druks_db.execute( +async def test_list_open_subjects_keeps_type_and_kind_partitions(client: TestClient, druks_db): + typed_note = await Note.create(body="two types") + note_run = await seed_run(druks_db, kind=Summarize.kind, subject=typed_note) + ticket_run = await seed_run(druks_db, kind=Summarize.kind, subject=typed_note) + await druks_db.execute( workflow_status.update() .where(workflow_status.c.workflow_uuid == ticket_run.id) .values( @@ -402,26 +411,26 @@ def test_list_open_subjects_keeps_type_and_kind_partitions(client: TestClient, d ) ) - multi_kind_note = Note.create(body="two kinds") - scan = seed_run(druks_db, kind="field_notes.scan", subject=multi_kind_note) - audit = seed_run(druks_db, kind="field_notes.audit", subject=multi_kind_note) + multi_kind_note = await Note.create(body="two kinds") + scan = await seed_run(druks_db, kind="field_notes.scan", subject=multi_kind_note) + audit = await seed_run(druks_db, kind="field_notes.audit", subject=multi_kind_note) - terminal_sibling_note = Note.create(body="terminal sibling") - failed_scan = seed_run( + terminal_sibling_note = await Note.create(body="terminal sibling") + failed_scan = await seed_run( druks_db, kind="field_notes.scan", subject=terminal_sibling_note, state="failed", ) failed_scan.created_at = datetime(2026, 1, 1, tzinfo=UTC) - finished_audit = seed_run( + finished_audit = await seed_run( druks_db, kind="field_notes.audit", subject=terminal_sibling_note, state="finished", ) finished_audit.created_at = failed_scan.created_at + timedelta(days=1) - druks_db.flush() + await druks_db.flush() body = client.get("/api/open-subjects").json() @@ -444,15 +453,15 @@ def test_list_open_subjects_keeps_type_and_kind_partitions(client: TestClient, d assert shared_id_types == {typed_note.subject_type, "ticket"} -def test_list_open_subjects_excludes_historical_runs(client: TestClient, druks_db): - note = Note.create(body="one open subject") +async def test_list_open_subjects_excludes_historical_runs(client: TestClient, druks_db): + note = await Note.create(body="one open subject") start = datetime(2026, 1, 1, tzinfo=UTC) for number in range(10): - historical = seed_run(druks_db, kind=Summarize.kind, subject=note, state="finished") + historical = await seed_run(druks_db, kind=Summarize.kind, subject=note, state="finished") historical.created_at = start + timedelta(seconds=number) - current = seed_run(druks_db, kind=Summarize.kind, subject=note) + current = await seed_run(druks_db, kind=Summarize.kind, subject=note) current.created_at = start + timedelta(seconds=10) - druks_db.flush() + await druks_db.flush() body = client.get("/api/open-subjects").json() @@ -461,24 +470,24 @@ def test_list_open_subjects_excludes_historical_runs(client: TestClient, druks_d ] == [current.id] -def test_list_open_subjects_caps_the_workflows(client: TestClient, druks_db): +async def test_list_open_subjects_caps_the_workflows(client: TestClient, druks_db): for number in range(51): - note = Note.create(body=f"open {number}") - seed_run(druks_db, kind=Summarize.kind, subject=note) + note = await Note.create(body=f"open {number}") + await seed_run(druks_db, kind=Summarize.kind, subject=note) body = client.get("/api/open-subjects").json() assert len(body["subjects"]) == 50 -def test_get_gate_then_answer_roundtrip(client: TestClient, druks_db, resume_spy): - note = Note.create(body="answer gate") - run = _park(druks_db, note) +async def test_get_gate_then_answer_roundtrip(client: TestClient, druks_db, resume_spy): + note = await Note.create(body="answer gate") + run = await _park(druks_db, note) view = client.get(f"/api/gates/{run.id}") assert view.status_code == 200 data = view.json() - assert data == services.get_gate(run.id).model_dump(mode="json", by_alias=True) + assert data == (await services.get_gate(run.id)).model_dump(mode="json", by_alias=True) answered = client.post( f"/api/gates/{run.id}/answer", @@ -489,14 +498,13 @@ def test_get_gate_then_answer_roundtrip(client: TestClient, druks_db, resume_spy assert resume_spy == [{"id": run.id, "action": "approve", "answers": {}, "note": "ship it"}] -def test_answer_gate_keys_empty_request_changes_on_ask_context( +async def test_answer_gate_keys_empty_request_changes_on_ask_context( client: TestClient, druks_db, resume_spy ): - critique_note = Note.create(body="critique-backed gate") - critique_run = _park(druks_db, critique_note, context="name the rollback boundary") - critique_parked_at = services.get_gate(critique_run.id).model_dump(mode="json", by_alias=True)[ - "parkedAt" - ] + critique_note = await Note.create(body="critique-backed gate") + critique_run = await _park(druks_db, critique_note, context="name the rollback boundary") + critique_gate = await services.get_gate(critique_run.id) + critique_parked_at = critique_gate.model_dump(mode="json", by_alias=True)["parkedAt"] answered = client.post( f"/api/gates/{critique_run.id}/answer", @@ -523,9 +531,9 @@ def test_answer_gate_keys_empty_request_changes_on_ask_context( } ] - contextless_note = Note.create(body="contextless gate") - contextless_run = _park(druks_db, contextless_note) - contextless_parked_at = services.get_gate(contextless_run.id).model_dump( + contextless_note = await Note.create(body="contextless gate") + contextless_run = await _park(druks_db, contextless_note) + contextless_parked_at = (await services.get_gate(contextless_run.id)).model_dump( mode="json", by_alias=True )["parkedAt"] @@ -548,15 +556,15 @@ def test_answer_gate_keys_empty_request_changes_on_ask_context( assert len(resume_spy) == 1 -def test_answer_gate_reads_already_answered_off_the_receipt( +async def test_answer_gate_reads_already_answered_off_the_receipt( client: TestClient, druks_db, resume_spy ): - note = Note.create(body="answered gate") + note = await Note.create(body="answered gate") parked_at = datetime.now(UTC) - run = seed_run(druks_db, kind=Summarize.kind, subject=note) + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) run.input_requested_at = parked_at run.answer_parked_at = parked_at - druks_db.flush() + await druks_db.flush() response = client.post( f"/api/gates/{run.id}/answer", @@ -568,9 +576,9 @@ def test_answer_gate_reads_already_answered_off_the_receipt( assert resume_spy == [] -def test_answer_gate_requires_an_aware_parked_at(client: TestClient, druks_db): - note = Note.create(body="naive parked timestamp") - run = _park(druks_db, note) +async def test_answer_gate_requires_an_aware_parked_at(client: TestClient, druks_db): + note = await Note.create(body="naive parked timestamp") + run = await _park(druks_db, note) naive = client.post( f"/api/gates/{run.id}/answer", @@ -580,14 +588,14 @@ def test_answer_gate_requires_an_aware_parked_at(client: TestClient, druks_db): assert naive.status_code == 422 # Pydantic's, not the agent taxonomy -def test_cancel_run_route(client: TestClient, druks_db): - note = Note.create(body="cancelled note") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) +async def test_cancel_run_route(client: TestClient, druks_db): + note = await Note.create(body="cancelled note") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) # Parked, so the cancel must clear the gate — and never write the receipt. run.input_gate = "review" run.input_request = {"presentation": "in_app", "questions": []} run.input_requested_at = run.utc_now() - druks_db.flush() + await druks_db.flush() unbounded = client.post(f"/api/runs/{run.id}/cancel", json={"reason": "r" * 501}) assert unbounded.status_code == 422 @@ -598,8 +606,8 @@ def test_cancel_run_route(client: TestClient, druks_db): assert cancelled.status_code == 200 assert cancelled.json() == {"run": run.id, "result": "cancelled"} - druks_db.expire_all() - run = druks_db.get(type(run), run.id) + druks_db.expunge_all() + run = await druks_db.get(type(run), run.id) assert not run.answer_parked_at assert not run.input_gate assert run.failure == "wrong branch" @@ -609,10 +617,10 @@ def test_cancel_run_route(client: TestClient, druks_db): assert again.json()["result"] == "already_cancelled" -def test_transcript_route_matches_the_read_machinery(client: TestClient, druks_db): - note = Note.create(body="transcript route") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", status="running") +async def test_transcript_route_matches_the_read_machinery(client: TestClient, druks_db): + note = await Note.create(body="transcript route") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", status="running") call_dir = call.call_dir call_dir.mkdir(parents=True, exist_ok=True) (call_dir / "stdout.jsonl").write_bytes(b"hello " + "é".encode() + b" transcript") @@ -628,17 +636,17 @@ def test_transcript_route_matches_the_read_machinery(client: TestClient, druks_d assert response.json()["text"] == "hello �" -def test_resume_route_contract_is_preserved(client: TestClient, druks_db, resume_spy): +async def test_resume_route_contract_is_preserved(client: TestClient, druks_db, resume_spy): unknown = client.post("/api/runs/no-such-run/resume", json={"control": "approve"}) assert unknown.status_code == 404 - idle_note = Note.create(body="idle run") - idle = seed_run(druks_db, kind=Summarize.kind, subject=idle_note) + idle_note = await Note.create(body="idle run") + idle = await seed_run(druks_db, kind=Summarize.kind, subject=idle_note) not_waiting = client.post(f"/api/runs/{idle.id}/resume", json={"control": "approve"}) assert not_waiting.status_code == 409 - parked_note = Note.create(body="parked run") - run = _park(druks_db, parked_note) + parked_note = await Note.create(body="parked run") + run = await _park(druks_db, parked_note) bad_control = client.post(f"/api/runs/{run.id}/resume", json={"control": "merge"}) assert bad_control.status_code == 422 assert resume_spy == [] @@ -666,15 +674,15 @@ def test_resume_route_contract_is_preserved(client: TestClient, druks_db, resume run.answer_parked_at = run.input_requested_at run.input_gate = None run.input_request = None - druks_db.flush() + await druks_db.flush() late = client.post(f"/api/runs/{run.id}/resume", json={"control": "approve"}) assert late.status_code == 409 assert len(resume_spy) == 1 -def test_usage_agent_route_matches_the_service(client: TestClient, druks_db, account): - note = Note.create(body="usage route") - run = seed_run( +async def test_usage_agent_route_matches_the_service(client: TestClient, druks_db, account): + note = await Note.create(body="usage route") + run = await seed_run( druks_db, kind=Summarize.kind, subject=note, @@ -693,12 +701,12 @@ def test_usage_agent_route_matches_the_service(client: TestClient, druks_db, acc cost_metadata={"total_tokens": 500}, ) ) - druks_db.flush() + await druks_db.flush() response = client.get("/api/usage/summary") assert response.status_code == 200 body = response.json() - assert body == services.get_usage(account).model_dump(mode="json", by_alias=True) + assert body == (await services.get_usage(account)).model_dump(mode="json", by_alias=True) assert len(response.content) <= 4 * 1024 today = client.get("/api/usage/today").json() diff --git a/backend/tests/test_agent_services.py b/backend/tests/test_agent_services.py index 9e24d2d7..4d562ad0 100644 --- a/backend/tests/test_agent_services.py +++ b/backend/tests/test_agent_services.py @@ -27,8 +27,8 @@ def _data_dir(tmp_path, monkeypatch): @pytest.fixture -def account(druks_db): - return Account.get_or_create("op@example.com") +async def account(druks_db): + return await Account.get_or_create("op@example.com") @pytest.fixture @@ -50,8 +50,8 @@ def _in_app_ask(questions=()): } -def _park(druks_db, note, *, ask=None): - run = seed_note_run( +async def _park(druks_db, note, *, ask=None): + run = await seed_note_run( druks_db, note=note, state="parked", @@ -59,7 +59,7 @@ def _park(druks_db, note, *, ask=None): input_request=ask if ask is not None else _in_app_ask(), ) run.input_requested_at = datetime.now(UTC) - druks_db.flush() + await druks_db.flush() return run @@ -102,12 +102,12 @@ def test_read_slice_missing_file_is_an_empty_eof(tmp_path: Path): # ---- gates ---------------------------------------------------------------- -def test_get_gate_returns_the_ask_and_parked_at(druks_db): - item = make_test_note() +async def test_get_gate_returns_the_ask_and_parked_at(druks_db): + item = await make_test_note() question = {"id": "q1", "prompt": "Which db?", "options": [{"id": "pg", "label": "Postgres"}]} - run = _park(druks_db, item, ask=_in_app_ask([question])) + run = await _park(druks_db, item, ask=_in_app_ask([question])) - view = services.get_gate(run.id) + view = await services.get_gate(run.id) assert view.run == run.id assert view.gate == "review" @@ -116,15 +116,15 @@ def test_get_gate_returns_the_ask_and_parked_at(druks_db): assert view.ask["questions"][0]["prompt"] == "Which db?" -def test_get_gate_serves_the_artifact(druks_db): - item = make_test_note() - run = _park(druks_db, item) - call = seed_call(druks_db, run, "generate_plan") - Artifact.record( +async def test_get_gate_serves_the_artifact(druks_db): + item = await make_test_note() + run = await _park(druks_db, item) + call = await seed_call(druks_db, run, "generate_plan") + await Artifact.record( call_dir=call.call_dir, call_id=call.id, kind="markdown", title="Plan", content="x" * 10240 ) - view = services.get_gate(run.id) + view = await services.get_gate(run.id) assert view.artifact is not None assert view.artifact.call_id == call.id @@ -132,23 +132,23 @@ def test_get_gate_serves_the_artifact(druks_db): assert len(view.artifact.content.encode()) <= 4096 -def test_get_gate_refuses_when_not_parked_or_external(druks_db): +async def test_get_gate_refuses_when_not_parked_or_external(druks_db): with pytest.raises(RunNotFound): - services.get_gate("no-such-run") + await services.get_gate("no-such-run") - item = make_test_note() - running = seed_note_run(druks_db, note=item, state="running") + item = await make_test_note() + running = await seed_note_run(druks_db, note=item, state="running") with pytest.raises(exceptions.GateNotOpen): - services.get_gate(running.id) + await services.get_gate(running.id) - external_item = make_test_note() - external = _park( + external_item = await make_test_note() + external = await _park( druks_db, external_item, ask={"presentation": "external", "label": "Answer on the ticket"}, ) with pytest.raises(exceptions.GateNotAnswerable): - services.get_gate(external.id) + await services.get_gate(external.id) # The answered and already-answered happy paths are pinned at both doors — @@ -162,15 +162,15 @@ async def test_answer_gate_error_taxonomy(druks_db, resume_spy): "no-such-run", parked_at=datetime.now(UTC), control="approve", answers={}, note="" ) - item = make_test_note() - finished = seed_note_run(druks_db, note=item, state="finished") + item = await make_test_note() + finished = await seed_note_run(druks_db, note=item, state="finished") with pytest.raises(exceptions.GateNotOpen): await services.answer_gate( finished.id, parked_at=datetime.now(UTC), control="approve", answers={}, note="" ) - parked_item = make_test_note() - run = _park(druks_db, parked_item) + parked_item = await make_test_note() + run = await _park(druks_db, parked_item) with pytest.raises(exceptions.GateRoundStale): await services.answer_gate( run.id, @@ -184,8 +184,8 @@ async def test_answer_gate_error_taxonomy(druks_db, resume_spy): run.id, parked_at=run.input_requested_at, control="merge", answers={}, note="" ) - external_item = make_test_note() - external = _park( + external_item = await make_test_note() + external = await _park( druks_db, external_item, ask={"presentation": "external", "label": "Answer on the ticket"}, @@ -204,30 +204,30 @@ async def test_answer_gate_error_taxonomy(druks_db, resume_spy): # ---- agent calls ---------------------------------------------------------- -def test_agent_call_get_returns_the_call_or_raises(druks_db): - run = seed_note_run(druks_db) - call = seed_call(druks_db, run, "summarize") +async def test_agent_call_get_returns_the_call_or_raises(druks_db): + run = await seed_note_run(druks_db) + call = await seed_call(druks_db, run, "summarize") - assert AgentCall.get(call.id) == call + assert (await AgentCall.get(call.id)).id == call.id with pytest.raises(AgentCallNotFound) as error: - AgentCall.get("missing") + await AgentCall.get("missing") assert str(error.value) == "No agent call missing." -def test_get_agent_call_serves_bounded_tails(druks_db): +async def test_get_agent_call_serves_bounded_tails(druks_db): from conftest import finish_agent_run, seed_note_agent_run - call = seed_note_agent_run() + call = await seed_note_agent_run() call_dir = call.call_dir call_dir.mkdir(parents=True, exist_ok=True) (call_dir / "stdout.jsonl").write_bytes(b"s" * 20480) (call_dir / "stderr.log").write_bytes(b"e" * 10240) - finish_agent_run(call, last_error="boom " * 100) - Artifact.record( + await finish_agent_run(call, last_error="boom " * 100) + await Artifact.record( call_dir=call_dir, call_id=call.id, kind="markdown", title="Out", content="a" * 10240 ) - detail = services.get_agent_call(call.id) + detail = await services.get_agent_call(call.id) assert detail.run == call.run_id assert detail.call.id == call.id @@ -238,22 +238,22 @@ def test_get_agent_call_serves_bounded_tails(druks_db): assert detail.artifact.content == "a" * 4096 with pytest.raises(AgentCallNotFound): - services.get_agent_call("no-such-call") + await services.get_agent_call("no-such-call") -def test_get_agent_call_without_files_reads_empty(druks_db): +async def test_get_agent_call_without_files_reads_empty(druks_db): from conftest import seed_note_agent_run - call = seed_note_agent_run() + call = await seed_note_agent_run() - detail = services.get_agent_call(call.id) + detail = await services.get_agent_call(call.id) assert detail.transcript == "" assert detail.stderr == "" assert detail.artifact is None -def test_artifact_content_omits_an_unknown_call(druks_db): +async def test_artifact_content_omits_an_unknown_call(druks_db): artifact = Artifact( agent_call_id="missing", kind="markdown", @@ -261,27 +261,27 @@ def test_artifact_content_omits_an_unknown_call(druks_db): path="artifact.md", ) - assert services._artifact_content(artifact) is None + assert await services._artifact_content(artifact) is None # ---- cancel --------------------------------------------------------------- async def test_cancel_run_paths(druks_db): - item = make_test_note() - run = seed_note_run(druks_db, note=item, state="running") + item = await make_test_note() + run = await seed_note_run(druks_db, note=item, state="running") result = await runs.cancel_run(run.id, reason="stuck") assert result.result == "cancelled" - druks_db.expire_all() - assert Run.get(run.id).state == "cancelled" - assert Run.get(run.id).failure == "stuck" + druks_db.expunge_all() + assert (await Run.get(run.id)).state == "cancelled" + assert (await Run.get(run.id)).failure == "stuck" again = await runs.cancel_run(run.id, reason="stuck") assert again.result == "already_cancelled" - finished_item = make_test_note() - finished = seed_note_run(druks_db, note=finished_item, state="finished") + finished_item = await make_test_note() + finished = await seed_note_run(druks_db, note=finished_item, state="finished") with pytest.raises(RunNotActive): await runs.cancel_run(finished.id, reason="late") @@ -290,8 +290,8 @@ async def test_cancel_run_paths(druks_db): async def test_run_retry_forks_from_the_failed_step(druks_db, monkeypatch): - item = make_test_note() - run = seed_note_run(druks_db, note=item, state="failed") + item = await make_test_note() + run = await seed_note_run(druks_db, note=item, state="failed") retried_run_id = "retried-run" steps = [ {"function_id": 2, "error": None}, @@ -302,7 +302,7 @@ async def test_run_retry_forks_from_the_failed_step(druks_db, monkeypatch): events = [] async def _fork(workflow_id, start_step, *, queue_name): - seed_dbos_status( + await seed_dbos_status( druks_db, retried_run_id, "scheduled", @@ -323,8 +323,8 @@ async def _publish(event, **facts): assert result == retried_run_id list_steps.assert_awaited_once_with(run.id) fork.assert_awaited_once_with(run.id, 8, queue_name=run_queue.name) - druks_db.expire_all() - retried = Run.get(retried_run_id) + druks_db.expunge_all() + retried = await Run.get(retried_run_id) assert retried.kind == run.kind assert retried.account_id == run.account_id assert retried.state == "scheduled" @@ -343,7 +343,7 @@ async def _publish(event, **facts): async def test_run_retry_restarts_at_step_one_without_a_failed_checkpoint(druks_db, monkeypatch): - run = seed_note_run(druks_db, state="failed") + run = await seed_note_run(druks_db, state="failed") list_steps = mock.AsyncMock(return_value=[{"function_id": 2, "error": None}]) fork = mock.AsyncMock(return_value=SimpleNamespace(workflow_id="clean-retry")) monkeypatch.setattr("dbos.DBOS.list_workflow_steps_async", list_steps) @@ -356,7 +356,7 @@ async def test_run_retry_restarts_at_step_one_without_a_failed_checkpoint(druks_ async def test_retry_run_refuses_a_non_failed_run(druks_db, monkeypatch): - run = seed_note_run(druks_db, state="finished") + run = await seed_note_run(druks_db, state="finished") retry = mock.AsyncMock() monkeypatch.setattr(Run, "retry", retry) @@ -369,9 +369,9 @@ async def test_retry_run_refuses_a_non_failed_run(druks_db, monkeypatch): async def test_retry_run_refuses_a_busy_subject(druks_db, monkeypatch): - item = make_test_note() - failed = seed_note_run(druks_db, note=item, state="failed") - active = seed_note_run(druks_db, note=item, state="running") + item = await make_test_note() + failed = await seed_note_run(druks_db, note=item, state="failed") + active = await seed_note_run(druks_db, note=item, state="running") retry = mock.AsyncMock() monkeypatch.setattr(Run, "retry", retry) @@ -384,7 +384,7 @@ async def test_retry_run_refuses_a_busy_subject(druks_db, monkeypatch): async def test_retry_run_retries_a_failed_run(druks_db, monkeypatch): - run = seed_note_run(druks_db, state="failed") + run = await seed_note_run(druks_db, state="failed") retry = mock.AsyncMock(return_value="retried-run") monkeypatch.setattr(Run, "retry", retry) @@ -405,13 +405,13 @@ async def test_retry_run_refuses_a_missing_run(druks_db): # ---- usage ---------------------------------------------------------------- -def test_get_usage_is_a_bounded_pure_read(druks_db, account): +async def test_get_usage_is_a_bounded_pure_read(druks_db, account): from druks.durable.models import AgentCall from druks.testing import seed_run from druks_field_notes.workflows import Summarize now = datetime.now(UTC) - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-usage") + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-usage") for index in range(30): druks_db.add( AgentCall( @@ -427,7 +427,7 @@ def test_get_usage_is_a_bounded_pure_read(druks_db, account): ) ) for tick in range(40): - UsageScrape( + await UsageScrape( harness="claude", account_id=account.id, scraped_at=now - timedelta(minutes=5 * tick), @@ -448,7 +448,8 @@ def test_get_usage_is_a_bounded_pure_read(druks_db, account): ], ).save() - usage = services.get_usage(account) + await druks_db.flush() + usage = await services.get_usage(account) assert usage.runs_today == 30 assert usage.spend_today_usd == pytest.approx(15.0) @@ -465,13 +466,13 @@ def test_get_usage_is_a_bounded_pure_read(druks_db, account): assert len(usage.model_dump_json(by_alias=True).encode()) <= 4 * 1024 -def test_get_usage_only_counts_the_callers_spend(druks_db, account): +async def test_get_usage_only_counts_the_callers_spend(druks_db, account): from druks.durable.models import AgentCall from druks.testing import seed_run from druks_field_notes.workflows import Summarize - other = Account.get_or_create("other@example.com") - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-usage-other") + other = await Account.get_or_create("other@example.com") + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-usage-other") druks_db.add( AgentCall( run_id=run.id, @@ -484,9 +485,9 @@ def test_get_usage_only_counts_the_callers_spend(druks_db, account): cost_usd=9.0, ) ) - druks_db.flush() + await druks_db.flush() - usage = services.get_usage(account) + usage = await services.get_usage(account) assert usage.runs_today == 0 assert usage.spend_today_usd == 0.0 diff --git a/backend/tests/test_agents.py b/backend/tests/test_agents.py index c1193aa3..c2ab8f96 100644 --- a/backend/tests/test_agents.py +++ b/backend/tests/test_agents.py @@ -21,7 +21,7 @@ class DummyOutput(agents.AgentOutput): ) -def test_get_timeout_caps_at_the_sandbox_lease_max(druks_db): +async def test_get_timeout_caps_at_the_sandbox_lease_max(druks_db): """A resolved timeout over the sandbox-lease max is clamped; a shorter one passes through.""" from druks.sandbox.constants import MAX_AGENT_TIMEOUT_SECONDS @@ -40,27 +40,27 @@ def test_get_timeout_caps_at_the_sandbox_lease_max(druks_db): timeout=600, ) - assert over.get_timeout() == MAX_AGENT_TIMEOUT_SECONDS - assert under.get_timeout() == 600 + assert await over.get_timeout() == MAX_AGENT_TIMEOUT_SECONDS + assert await under.get_timeout() == 600 @pytest.fixture(autouse=True) -def _seed_run_for_record(druks_db): +async def _seed_run_for_record(druks_db): # An agent call records an AgentCall, which FKs to its run. from druks.testing import seed_run from druks_field_notes.workflows import Summarize - seed_run(druks_db, kind=Summarize.kind, run_id="wf-9") + await seed_run(druks_db, kind=Summarize.kind, run_id="wf-9") @pytest.fixture(autouse=True) -def _connected_claude(druks_db): +async def _connected_claude(druks_db): # A run refuses to dispatch on an unconnected harness; the runtime tests # here resolve to claude models, so connect it once. from conftest import connect_harness from druks.harnesses.claude import ClaudeHarness - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "test-token"}}) + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "test-token"}}) def _patch_runtime(monkeypatch, tmp_path, payload): @@ -140,8 +140,10 @@ async def test_run_refuses_unconnected_harness(druks_db, tmp_path, monkeypatch, from druks.harnesses.exceptions import HarnessNotConnectedError from druks.harnesses.models import HarnessConnection - HarnessConnection.get_for_account( - "claude", Account.get_for_username("op@example.com").id + await ( + await HarnessConnection.get_for_account( + "claude", (await Account.get_for_username("op@example.com")).id + ) ).delete() sandbox = _patch_runtime(monkeypatch, tmp_path, {"ok": True}) _patch_ephemeral(monkeypatch, sandbox) @@ -244,7 +246,7 @@ async def test_running_call_visible_then_finished(druks_db, tmp_path, monkeypatc during: dict[str, object] = {} async def _run_agent(*, call_id, **_kwargs): - row = AgentCall.get(call_id) + row = await AgentCall.get(call_id) during["status"] = row.status during["host"] = row.sandbox_host_id return make_agent_result({"ok": True}, agent="dummy") @@ -255,7 +257,7 @@ async def _run_agent(*, call_id, **_kwargs): await DUMMY_AGENT._run(workflow_id="wf-9") assert during == {"status": "running", "host": "host-test"} - [call] = AgentCall.list_for_run("wf-9") + [call] = await AgentCall.list_for_run("wf-9") assert call.status == "succeeded" assert call.sandbox_host_id == "host-test" assert call.finished_at is not None @@ -276,7 +278,7 @@ async def boom(self, *, idempotency_key=None, **_kwargs): with pytest.raises(RuntimeError, match="no capacity"): await DUMMY_AGENT._run(workflow_id="wf-9") - assert AgentCall.list_for_run("wf-9") == [] + assert await AgentCall.list_for_run("wf-9") == [] async def test_crash_after_start_fails_the_call(druks_db, tmp_path, monkeypatch, current_run): @@ -293,7 +295,7 @@ async def _boom(**_kwargs): with pytest.raises(RuntimeError, match="kaboom"): await DUMMY_AGENT._run(workflow_id="wf-9") - [call] = AgentCall.list_for_run("wf-9") + [call] = await AgentCall.list_for_run("wf-9") assert call.status == "failed" assert "kaboom" in call.last_error @@ -319,7 +321,7 @@ async def _run_agent(**_kwargs): await DUMMY_AGENT._run(workflow_id="wf-9") assert excinfo.value is overloaded - [call] = AgentCall.list_for_run("wf-9") + [call] = await AgentCall.list_for_run("wf-9") assert call.status == "failed" assert call.failure_code == "overloaded" assert call.last_error == ( @@ -355,7 +357,7 @@ async def test_body_level_overload_retries_as_separate_durable_attempts( assert [awaited.args[0] for awaited in sleep.await_args_list] == [285.0, 945.0] assert [options["name"] for options in checkpoints] == ["test.agent.dummy"] * 3 assert current_run._reap_run.await_count == 2 - calls = AgentCall.list_for_run("wf-9") + calls = await AgentCall.list_for_run("wf-9") assert len(calls) == 3 assert sum(call.status == "failed" for call in calls) == 2 assert sum(call.status == "succeeded" for call in calls) == 1 @@ -385,11 +387,18 @@ async def test_body_level_first_byte_retries_immediately_then_reraises( assert excinfo.value.code == "first_byte" assert [awaited.args[0] for awaited in sleep.await_args_list] == [0.0, 0.0] current_run._reap_run.assert_not_awaited() - calls = AgentCall.list_for_run("wf-9") + calls = await AgentCall.list_for_run("wf-9") assert len(calls) == 3 assert all(call.status == "failed" for call in calls) +def _async_scrape(make): + async def latest_for(_cls, _harness, _account_id): + return make() + + return latest_for + + async def test_body_level_quota_waits_for_the_reset_once( druks_db, tmp_path, monkeypatch, current_run, _inline_agent_steps ): @@ -409,20 +418,22 @@ async def test_body_level_quota_waits_for_the_reset_once( agents.UsageScrape, "latest_for", classmethod( - lambda _cls, _harness, _account_id: UsageScrape( - five_hour_resets_at=now + timedelta(hours=2), - weeks=[ - { - "percent_left": 0, - "resets_at": (now + timedelta(hours=1)).isoformat(), - "model": "Fable", - }, - { - "percent_left": 20, - "resets_at": (now + timedelta(days=1)).isoformat(), - "model": None, - }, - ], + _async_scrape( + lambda: UsageScrape( + five_hour_resets_at=now + timedelta(hours=2), + weeks=[ + { + "percent_left": 0, + "resets_at": (now + timedelta(hours=1)).isoformat(), + "model": "Fable", + }, + { + "percent_left": 20, + "resets_at": (now + timedelta(days=1)).isoformat(), + "model": None, + }, + ], + ) ) ), ) @@ -443,7 +454,7 @@ async def test_body_level_quota_waits_for_the_reset_once( "test.agent.dummy.retry_wait", "test.agent.dummy", ] - calls = AgentCall.list_for_run("wf-9") + calls = await AgentCall.list_for_run("wf-9") assert len(calls) == 2 assert all(call.failure_code == "rate_limited" for call in calls) @@ -467,15 +478,17 @@ async def test_body_level_quota_reset_over_six_hours_reraises_without_sleeping( agents.UsageScrape, "latest_for", classmethod( - lambda _cls, _harness, _account_id: UsageScrape( - five_hour_resets_at=now + timedelta(hours=7), - weeks=[ - { - "percent_left": 0, - "resets_at": (now + timedelta(days=1)).isoformat(), - "model": None, - } - ], + _async_scrape( + lambda: UsageScrape( + five_hour_resets_at=now + timedelta(hours=7), + weeks=[ + { + "percent_left": 0, + "resets_at": (now + timedelta(days=1)).isoformat(), + "model": None, + } + ], + ) ) ), ) @@ -491,7 +504,7 @@ async def test_body_level_quota_reset_over_six_hours_reraises_without_sleeping( sleep.assert_not_awaited() jitter.assert_not_called() current_run._reap_run.assert_not_awaited() - [call] = AgentCall.list_for_run("wf-9") + [call] = await AgentCall.list_for_run("wf-9") assert call.failure_code == "usage_limit" @@ -520,7 +533,7 @@ async def test_body_level_never_retry_errors_run_once( sleep.assert_not_awaited() current_run._reap_run.assert_not_awaited() sandbox.run_agent.assert_awaited_once() - assert len(AgentCall.list_for_run("wf-9")) == 1 + assert len(await AgentCall.list_for_run("wf-9")) == 1 async def test_in_step_transient_retry_uses_asyncio_sleep(monkeypatch, current_run): @@ -766,7 +779,7 @@ async def test_recovery_supersedes_the_orphaned_running_call(druks_db): from druks.durable.engine import _step_engine engine = _step_engine() - AgentCall.start( + await AgentCall.start( engine, call_id="a", run_id="wf-9", @@ -775,7 +788,7 @@ async def test_recovery_supersedes_the_orphaned_running_call(druks_db): host_id="h", account_id="system", ) - AgentCall.start( + await AgentCall.start( engine, call_id="b", run_id="wf-9", @@ -785,7 +798,7 @@ async def test_recovery_supersedes_the_orphaned_running_call(druks_db): account_id="system", ) - by_id = {call.id: call for call in AgentCall.list_for_run("wf-9")} + by_id = {call.id: call for call in await AgentCall.list_for_run("wf-9")} assert by_id["a"].status == "abandoned" assert by_id["a"].finished_at is not None assert by_id["b"].status == "running" diff --git a/backend/tests/test_api_runs.py b/backend/tests/test_api_runs.py index 7202bbc7..f45e30b9 100644 --- a/backend/tests/test_api_runs.py +++ b/backend/tests/test_api_runs.py @@ -24,16 +24,16 @@ def client(tmp_path: Path, druks_db, monkeypatch): yield client -def _seed_call_id( +async def _seed_call_id( druks_db, *, finished: bool = True, run_state: str = "running", ) -> str: - note = Note.create(body="agent transcript") - run = seed_run(druks_db, kind=Summarize.kind, subject=note, state=run_state) + note = await Note.create(body="agent transcript") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note, state=run_state) status = "succeeded" if finished else "running" - call = seed_call(druks_db, run, "summarize", status=status) + call = await seed_call(druks_db, run, "summarize", status=status) # The harness writes every file for a call into run-//. call_dir = call.call_dir @@ -46,11 +46,11 @@ def _seed_call_id( return call.id -def test_list_files_inventories_call_artifacts( +async def test_list_files_inventories_call_artifacts( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) response = client.get(f"/api/field_notes/transcripts/{call_id}/files") @@ -64,9 +64,9 @@ def test_list_files_inventories_call_artifacts( assert files["metadata"] is not None -def test_get_agent_call_files_raises_for_unknown_call(druks_db): +async def test_get_agent_call_files_raises_for_unknown_call(druks_db): with pytest.raises(AgentCallNotFound): - get_agent_call_files("missing") + await get_agent_call_files("missing") def test_transcript_unknown_call_returns_unified_404(client: TestClient, druks_db): @@ -105,11 +105,11 @@ def test_file_download_unknown_call_returns_unified_404(client: TestClient, druk } -def test_transcript_range_fetch_paginates( +async def test_transcript_range_fetch_paginates( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) first = client.get( f"/api/field_notes/transcripts/{call_id}", @@ -133,11 +133,11 @@ def test_transcript_range_fetch_paginates( assert data["eof"] is True -def test_transcript_of_a_running_call_is_never_cached( +async def test_transcript_of_a_running_call_is_never_cached( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db, finished=False) + call_id = await _seed_call_id(druks_db, finished=False) response = client.get( f"/api/field_notes/transcripts/{call_id}", @@ -148,13 +148,13 @@ def test_transcript_of_a_running_call_is_never_cached( assert response.headers["cache-control"] == "no-store" -def test_transcript_of_an_abandoned_call_is_cached_immutably( +async def test_transcript_of_an_abandoned_call_is_cached_immutably( client: TestClient, druks_db, ): # The call never wrote finished_at, but its run died — nothing will append # to that log again, so the chunk is as permanent as a finished call's. - call_id = _seed_call_id(druks_db, finished=False, run_state="failed") + call_id = await _seed_call_id(druks_db, finished=False, run_state="failed") response = client.get( f"/api/field_notes/transcripts/{call_id}", @@ -165,13 +165,13 @@ def test_transcript_of_an_abandoned_call_is_cached_immutably( assert response.headers["cache-control"] == "public, max-age=31536000, immutable" -def test_transcript_missing_file_returns_eof( +async def test_transcript_missing_file_returns_eof( client: TestClient, druks_db, ): - note = Note.create(body="missing transcript") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", status="running") + note = await Note.create(body="missing transcript") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", status="running") response = client.get( f"/api/field_notes/transcripts/{call.id}", @@ -184,13 +184,13 @@ def test_transcript_missing_file_returns_eof( assert data["text"] == "" -def test_transcript_stream_emits_chunk_then_finishes( +async def test_transcript_stream_emits_chunk_then_finishes( client: TestClient, druks_db, ): # The terminal seeded run streams its stdout in one tick: a transcript.chunk # carrying the log, then agent_call.finished (which ends the SSE). - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) response = client.get( f"/api/field_notes/transcripts/{call_id}/stream", @@ -218,11 +218,11 @@ def test_transcript_stream_unknown_call_closes( assert response.text == "" -def test_get_file_serves_inventory_paths( +async def test_get_file_serves_inventory_paths( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) files = client.get(f"/api/field_notes/transcripts/{call_id}/files").json() # Compose the download URL the way the client does: the listing's own route @@ -234,11 +234,11 @@ def test_get_file_serves_inventory_paths( assert response.json() == {"ok": True} -def test_get_file_rejects_path_traversal( +async def test_get_file_rejects_path_traversal( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) response = client.get( f"/api/field_notes/transcripts/{call_id}/files/..%2F..%2Fetc%2Fpasswd", @@ -251,11 +251,11 @@ def test_get_file_rejects_path_traversal( } -def test_get_file_missing_returns_404( +async def test_get_file_missing_returns_404( client: TestClient, druks_db, ): - call_id = _seed_call_id(druks_db) + call_id = await _seed_call_id(druks_db) response = client.get( f"/api/field_notes/transcripts/{call_id}/files/nope.json", @@ -268,12 +268,12 @@ def test_get_file_missing_returns_404( } -def test_agent_call_artifact_layout(druks_db, tmp_path): +async def test_agent_call_artifact_layout(druks_db, tmp_path): # Layout sub-dir is the call id; the sandbox runner streams every run's stdout # to stdout.jsonl, so the layout is the same whichever harness ran the call. - note = Note.create(body="artifact layout") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", model="claude-opus-4-7", status="running") + note = await Note.create(body="artifact layout") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", model="claude-opus-4-7", status="running") sub = Path(call.artifact_dir) / call.id assert call.artifact_layout.transcript == sub / "stdout.jsonl" assert call.artifact_layout.stderr == sub / "stderr.log" diff --git a/backend/tests/test_api_settings.py b/backend/tests/test_api_settings.py index e5ffb330..da2c0c97 100644 --- a/backend/tests/test_api_settings.py +++ b/backend/tests/test_api_settings.py @@ -43,13 +43,13 @@ def test_harness_response_carries_connection_state(tmp_path: Path): assert "expiresAt" in claude -def test_harnesses_show_only_the_requesting_accounts_connection(tmp_path: Path, druks_db): +async def test_harnesses_show_only_the_requesting_accounts_connection(tmp_path: Path, druks_db): from conftest import connect_harness from druks.harnesses.claude import ClaudeHarness # The suite's identity gate stands in op@example.com; another account's # connection never shows on this card. - connect_harness( + await connect_harness( ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}, provider_email="someone-else@example.com", @@ -59,14 +59,14 @@ def test_harnesses_show_only_the_requesting_accounts_connection(tmp_path: Path, assert claude["connected"] is False -def test_harness_card_reports_identity(tmp_path: Path, druks_db): +async def test_harness_card_reports_identity(tmp_path: Path, druks_db): from druks.accounts.models import Account from druks.harnesses.models import HarnessConnection # The provider identity is display, never authority. - HarnessConnection.connect( + await HarnessConnection.connect( harness="claude", - account=Account.get_or_create("op@example.com"), + account=await Account.get_or_create("op@example.com"), payload={"claudeAiOauth": {"accessToken": "x"}}, expires_at=None, provider_email="seat@corp.com", @@ -78,15 +78,15 @@ def test_harness_card_reports_identity(tmp_path: Path, druks_db): assert claude["providerEmail"] == "seat@corp.com" -def test_harness_card_reads_expired_token_as_not_connected(tmp_path: Path, druks_db): +async def test_harness_card_reads_expired_token_as_not_connected(tmp_path: Path, druks_db): from datetime import UTC, datetime, timedelta from druks.accounts.models import Account from druks.harnesses.models import HarnessConnection - HarnessConnection.connect( + await HarnessConnection.connect( harness="claude", - account=Account.get_or_create("op@example.com"), + account=await Account.get_or_create("op@example.com"), payload={"claudeAiOauth": {"accessToken": "x"}}, expires_at=datetime.now(UTC) - timedelta(hours=1), provider_email="seat@corp.com", @@ -96,13 +96,13 @@ def test_harness_card_reads_expired_token_as_not_connected(tmp_path: Path, druks assert claude["connected"] is False -def test_disconnect_removes_only_the_requesting_accounts_connection(tmp_path: Path, druks_db): +async def test_disconnect_removes_only_the_requesting_accounts_connection(tmp_path: Path, druks_db): from conftest import connect_harness from druks.harnesses.claude import ClaudeHarness from druks.harnesses.models import HarnessConnection - mine = connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) - other = connect_harness( + mine = await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) + other = await connect_harness( ClaudeHarness, {"claudeAiOauth": {"accessToken": "y"}}, provider_email="someone-else@example.com", @@ -114,8 +114,8 @@ def test_disconnect_removes_only_the_requesting_accounts_connection(tmp_path: Pa assert response.json()["connected"] is False # The request deleted in its own task-scoped session; read past this # task's identity map for what actually persisted. - assert not HarnessConnection.reload(mine_id) - assert HarnessConnection.reload(other_id) + assert not await HarnessConnection.reload(mine_id) + assert await HarnessConnection.reload(other_id) def test_disconnect_without_a_connection_is_a_no_op(tmp_path: Path): @@ -126,7 +126,10 @@ def test_disconnect_without_a_connection_is_a_no_op(tmp_path: Path): def test_patch_settings_persists_valid_iana_zone(tmp_path: Path, monkeypatch): - monkeypatch.setattr("druks.user_settings.routes.apply_schedules", lambda: None) + async def _noop_schedules(): + return None + + monkeypatch.setattr("druks.user_settings.routes.apply_schedules", _noop_schedules) with _build_client(tmp_path) as client: patch = client.patch("/api/settings", json={"timezone": "Europe/Madrid"}) assert patch.status_code == 200 @@ -150,9 +153,11 @@ def test_timezone_change_reconciles_schedules(tmp_path: Path, monkeypatch): """Crons are evaluated in the operator's timezone, so changing it repoints the DBOS schedules now; re-asserting the same zone doesn't churn them.""" reconciled = [] - monkeypatch.setattr( - "druks.user_settings.routes.apply_schedules", lambda: reconciled.append(True) - ) + + async def record(): + reconciled.append(True) + + monkeypatch.setattr("druks.user_settings.routes.apply_schedules", record) with _build_client(tmp_path) as client: patch = client.patch("/api/settings", json={"timezone": "Europe/Madrid"}) assert patch.status_code == 200 @@ -274,7 +279,7 @@ def test_apps_surface_build_agents_and_workflow_defaults(tmp_path: Path): } -def test_app_secret_round_trip_encrypts_at_rest(tmp_path: Path): +async def test_app_secret_round_trip_encrypts_at_rest(tmp_path: Path): secret = "review-pem-value" app_id = "42424242" key = "app:review:private_key" @@ -291,18 +296,16 @@ def test_app_secret_round_trip_encrypts_at_rest(tmp_path: Path): }, ) stored = ( - db_session() - .execute( + await db_session().execute( text( "SELECT value, value IS NULL AS value_is_null, secret_value " "FROM settings_overrides WHERE key = :key" ), {"key": key}, ) - .one() - ) + ).one() read = client.get("/api/settings/apps") - resolved = Review.settings().private_key + resolved = (await Review.settings()).private_key assert written.status_code == 200 assert read.status_code == 200 @@ -325,15 +328,15 @@ def test_app_secret_round_trip_encrypts_at_rest(tmp_path: Path): assert fields["app_id"]["secretSet"] is True -def test_app_secret_plaintext_row_is_unset_until_resaved(tmp_path: Path): +async def test_app_secret_plaintext_row_is_unset_until_resaved(tmp_path: Path): secret = "legacy-plaintext-secret" key = "app:review:private_key" db_session().add(SettingsOverride(key=key, value=secret)) - db_session().flush() + await db_session().flush() with _build_client(tmp_path) as client: initial = _review_app(client) - resolved_initial = Review.settings().private_key + resolved_initial = (await Review.settings()).private_key saved = client.patch( "/api/settings/apps", json={ @@ -346,16 +349,14 @@ def test_app_secret_plaintext_row_is_unset_until_resaved(tmp_path: Path): }, ) stored = ( - db_session() - .execute( + await db_session().execute( text( "SELECT value, value IS NULL AS value_is_null, secret_value " "FROM settings_overrides WHERE key = :key" ), {"key": key}, ) - .one() - ) + ).one() initial_field = next( setting for setting in initial["settings"] if setting["name"] == "private_key" @@ -369,7 +370,7 @@ def test_app_secret_plaintext_row_is_unset_until_resaved(tmp_path: Path): assert secret.encode() not in stored.secret_value -def test_app_non_secret_setting_stays_in_value(tmp_path: Path): +async def test_app_non_secret_setting_stays_in_value(tmp_path: Path): status = "Agent Queue" key = "app:ship:linear_trigger_status" @@ -379,13 +380,11 @@ def test_app_non_secret_setting_stays_in_value(tmp_path: Path): json={"appSettings": {"ship": {"linear_trigger_status": status}}}, ) stored = ( - db_session() - .execute( + await db_session().execute( text("SELECT value, secret_value FROM settings_overrides WHERE key = :key"), {"key": key}, ) - .one() - ) + ).one() ship = _ship_app(client) field = next( @@ -394,7 +393,7 @@ def test_app_non_secret_setting_stays_in_value(tmp_path: Path): assert written.status_code == 200 assert stored.value == status assert stored.secret_value == b"" - assert Ship.settings().linear_trigger_status == status + assert (await Ship.settings()).linear_trigger_status == status assert field["value"] == status assert field["overridden"] is True @@ -403,9 +402,11 @@ def test_incoherent_app_save_is_rejected_and_rolled_back_before_schedules( tmp_path: Path, monkeypatch ): reconciled = [] - monkeypatch.setattr( - "druks.user_settings.routes.apply_schedules", lambda: reconciled.append(True) - ) + + async def record(): + reconciled.append(True) + + monkeypatch.setattr("druks.user_settings.routes.apply_schedules", record) with _build_client(tmp_path) as client: response = client.patch( "/api/settings/apps", @@ -425,7 +426,7 @@ def test_incoherent_app_save_is_rejected_and_rolled_back_before_schedules( assert agents["generate_plan"]["model"] == "gpt-5.5" -def test_clearing_the_identity_deletes_its_overrides_and_stays_coherent(tmp_path: Path): +async def test_clearing_the_identity_deletes_its_overrides_and_stays_coherent(tmp_path: Path): key = "app:review:app_id" with _build_client(tmp_path) as client: @@ -445,19 +446,17 @@ def test_clearing_the_identity_deletes_its_overrides_and_stays_coherent(tmp_path json={"appSettings": {"review": {"app_id": None, "private_key": None}}}, ) stored = ( - db_session() - .execute( + await db_session().execute( text("SELECT 1 FROM settings_overrides WHERE key = :key"), {"key": key}, ) - .one_or_none() - ) + ).one_or_none() fields = _review_settings_fields(client) assert configured.status_code == 200 assert cleared.status_code == 200 assert stored is None - assert not Review.settings().app_id + assert not (await Review.settings()).app_id assert fields["app_id"]["secretSet"] is False assert fields["private_key"]["secretSet"] is False diff --git a/backend/tests/test_api_usage.py b/backend/tests/test_api_usage.py index 3c2cf087..2d9ec1cf 100644 --- a/backend/tests/test_api_usage.py +++ b/backend/tests/test_api_usage.py @@ -27,44 +27,46 @@ def client(app_settings: Settings): yield c -def _account_id() -> str: +async def _account_id() -> str: # The suite's auth gate stands in op@example.com (conftest override). - return Account.get_or_create("op@example.com").id + return (await Account.get_or_create("op@example.com")).id -def _seed(snapshots: list[UsageScrape]) -> None: +async def _seed(snapshots: list[UsageScrape]) -> None: # save() flushes onto the ambient per-test connection session (bound by the # _txn fixture), so the rows are visible to the request and roll back with # the test — no separate engine, no commit. Every snapshot belongs to the # viewing account unless a test stamps another owner. - viewer = _account_id() + viewer = await _account_id() for snap in snapshots: if not snap.account_id: snap.account_id = viewer - snap.save() + await snap.save() def _harness(body: dict, name: str) -> dict: return next(entry for entry in body["harnesses"] if entry["name"] == name) -def _seed_agent_call(druks_db, *, model: str = "gpt-5.5"): - note = Note.create(body="usage accounting") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - return seed_call(druks_db, run, "summarize", status="running", model=model) +async def _seed_agent_call(druks_db, *, model: str = "gpt-5.5"): + note = await Note.create(body="usage accounting") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + return await seed_call(druks_db, run, "summarize", status="running", model=model) -def test_usage_today_counts_calls_whose_model_isnt_a_current_harness(client, druks_db) -> None: +async def test_usage_today_counts_calls_whose_model_isnt_a_current_harness( + client, druks_db +) -> None: # Model ids churn on deploys (opus-4-7 → 4-8), so a call finished earlier today # can carry an id no harness claims any more. Money spent must not vanish from # the display — the sys-strip's total_run_spend_between counts every call, and # the two surfaces must quote the same number. Unclaimed models land in the # "unattributed" bucket the panel's grand total sums. - call = _seed_agent_call(druks_db, model="claude-opus-4-5") - call.account_id = _account_id() + call = await _seed_agent_call(druks_db, model="claude-opus-4-5") + call.account_id = await _account_id() call.finished_at = datetime.now(UTC) call.cost_usd = 2.5 - druks_db.flush() + await druks_db.flush() body = client.get("/api/usage/today").json() bucket = _harness(body, "unattributed") @@ -81,10 +83,10 @@ def test_get_usage_empty_returns_available_false(client) -> None: assert all(entry["available"] is False for entry in body["harnesses"]) -def test_get_usage_serializes_latest_per_harness(client, app_settings) -> None: +async def test_get_usage_serializes_latest_per_harness(client, app_settings) -> None: # Plant a snapshot for claude only — codex should still report # ``available=false`` rather than missing-key/404. - _seed( + await _seed( [ UsageScrape( harness="claude", @@ -120,8 +122,8 @@ def test_get_usage_serializes_latest_per_harness(client, app_settings) -> None: assert _harness(body, "codex")["available"] is False -def test_get_usage_flags_stale_after_24h(client, app_settings) -> None: - _seed( +async def test_get_usage_flags_stale_after_24h(client, app_settings) -> None: + await _seed( [ UsageScrape( harness="claude", @@ -136,11 +138,11 @@ def test_get_usage_flags_stale_after_24h(client, app_settings) -> None: assert _harness(body, "claude")["stale"] is True -def test_get_usage_exposes_unlimited_flag(client, app_settings) -> None: +async def test_get_usage_exposes_unlimited_flag(client, app_settings) -> None: # Codex business plan: scraper synthesizes permanently-full buckets # and marks the row unmetered so the UI can render "unmetered" # instead of a quota bar that never moves. - _seed( + await _seed( [ UsageScrape( harness="codex", @@ -159,7 +161,7 @@ def test_get_usage_exposes_unlimited_flag(client, app_settings) -> None: assert _harness(body, "claude")["unlimited"] is False -def test_usage_history_serializes_series_oldest_first(client, app_settings) -> None: +async def test_usage_history_serializes_series_oldest_first(client, app_settings) -> None: now = datetime.now(UTC) snaps = [ UsageScrape( @@ -191,7 +193,7 @@ def test_usage_history_serializes_series_oldest_first(client, app_settings) -> N snaps.append( UsageScrape(harness="claude", parse_ok=False, scraped_at=now - timedelta(minutes=5)) ) - _seed(snaps) + await _seed(snaps) body = client.get("/api/usage/history").json() @@ -213,11 +215,11 @@ def test_usage_history_serializes_series_oldest_first(client, app_settings) -> N assert _harness(body, "codex")["weeks"] == [] -def test_usage_today_aggregates_spend_and_tokens_by_provider( +async def test_usage_today_aggregates_spend_and_tokens_by_provider( client, app_settings, druks_db ) -> None: - codex_run = _seed_agent_call(druks_db, model="gpt-5.5") - codex_run.account_id = _account_id() + codex_run = await _seed_agent_call(druks_db, model="gpt-5.5") + codex_run.account_id = await _account_id() codex_run.cost_usd = 1.25 codex_run.cost_metadata = { "provider": "openai", @@ -227,8 +229,8 @@ def test_usage_today_aggregates_spend_and_tokens_by_provider( } codex_run.finished_at = datetime.now(UTC) - claude_run = _seed_agent_call(druks_db, model="claude-opus-4-7") - claude_run.account_id = _account_id() + claude_run = await _seed_agent_call(druks_db, model="claude-opus-4-7") + claude_run.account_id = await _account_id() claude_run.cost_usd = 2.5 claude_run.cost_metadata = { "provider": "anthropic", @@ -240,13 +242,13 @@ def test_usage_today_aggregates_spend_and_tokens_by_provider( claude_run.finished_at = datetime.now(UTC) # Finished yesterday — outside today's boundary, must not count. - old_run = _seed_agent_call(druks_db, model="gpt-5.5") + old_run = await _seed_agent_call(druks_db, model="gpt-5.5") old_run.cost_usd = 99.0 old_run.finished_at = datetime.now(UTC) - timedelta(days=2) # Still running — no cost yet, counted nowhere. - _seed_agent_call(druks_db, model="gpt-5.5") - druks_db.flush() + await _seed_agent_call(druks_db, model="gpt-5.5") + await druks_db.flush() body = client.get("/api/usage/today").json() @@ -266,10 +268,10 @@ def test_usage_today_aggregates_spend_and_tokens_by_provider( assert sum(claude["hours"]) == 2.5 -def test_usage_excludes_another_accounts_scrape(client, druks_db) -> None: +async def test_usage_excludes_another_accounts_scrape(client, druks_db) -> None: snap = UsageScrape(harness="claude", parse_ok=True, five_hour_percent_left=54) - snap.account_id = Account.get_or_create("other@example.com").id - snap.save() + snap.account_id = (await Account.get_or_create("other@example.com")).id + await snap.save() body = client.get("/api/usage").json() assert _harness(body, "claude")["available"] is False @@ -277,10 +279,10 @@ def test_usage_excludes_another_accounts_scrape(client, druks_db) -> None: assert _harness(history, "claude")["fiveHour"] == [] -def test_usage_reports_viewers_connection_identity(client, druks_db) -> None: - HarnessConnection.connect( +async def test_usage_reports_viewers_connection_identity(client, druks_db) -> None: + await HarnessConnection.connect( harness="claude", - account=Account.get_or_create("other@example.com"), + account=await Account.get_or_create("other@example.com"), payload={"claudeAiOauth": {"accessToken": "other"}}, expires_at=None, provider_email="other-seat@example.com", @@ -289,9 +291,9 @@ def test_usage_reports_viewers_connection_identity(client, druks_db) -> None: assert _harness(body, "claude")["connected"] is False assert _harness(body, "claude")["providerEmail"] is None - HarnessConnection.connect( + await HarnessConnection.connect( harness="claude", - account=Account.get_or_create("op@example.com"), + account=await Account.get_or_create("op@example.com"), payload={"claudeAiOauth": {"accessToken": "mine"}}, expires_at=None, provider_email="subscription@example.com", @@ -301,21 +303,21 @@ def test_usage_reports_viewers_connection_identity(client, druks_db) -> None: assert _harness(body, "claude")["providerEmail"] == "subscription@example.com" -def test_usage_today_counts_only_the_viewers_calls(client, druks_db) -> None: - mine = _seed_agent_call(druks_db, model="claude-opus-4-7") - mine.account_id = _account_id() +async def test_usage_today_counts_only_the_viewers_calls(client, druks_db) -> None: + mine = await _seed_agent_call(druks_db, model="claude-opus-4-7") + mine.account_id = await _account_id() mine.cost_usd = 2.0 mine.finished_at = datetime.now(UTC) - other = _seed_agent_call(druks_db, model="claude-opus-4-7") - other.account_id = Account.get_or_create("other@example.com").id + other = await _seed_agent_call(druks_db, model="claude-opus-4-7") + other.account_id = (await Account.get_or_create("other@example.com")).id other.cost_usd = 5.0 other.finished_at = datetime.now(UTC) - background = _seed_agent_call(druks_db, model="claude-opus-4-7") + background = await _seed_agent_call(druks_db, model="claude-opus-4-7") background.cost_usd = 9.0 background.finished_at = datetime.now(UTC) - druks_db.flush() + await druks_db.flush() body = client.get("/api/usage/today").json() assert _harness(body, "claude")["spendUsd"] == 2.0 @@ -338,9 +340,9 @@ async def fake(connection, *, now=None): return fake -def test_refresh_scrapes_only_the_viewers_connections(client, druks_db, monkeypatch) -> None: - viewer = connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) - connect_harness( +async def test_refresh_scrapes_only_the_viewers_connections(client, druks_db, monkeypatch) -> None: + viewer = await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) + await connect_harness( ClaudeHarness, {"claudeAiOauth": {"accessToken": "t2"}}, provider_email="other@example.com" ) fetched: list[str] = [] @@ -348,11 +350,11 @@ def test_refresh_scrapes_only_the_viewers_connections(client, druks_db, monkeypa assert client.post("/api/usage/refresh").status_code == 200 assert fetched == [viewer.account_id] - assert UsageScrape.latest_for("claude", viewer.account_id).five_hour_percent_left == 50 + assert (await UsageScrape.latest_for("claude", viewer.account_id)).five_hour_percent_left == 50 -def test_refresh_floors_repeat_scrapes(client, druks_db, monkeypatch) -> None: - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) +async def test_refresh_floors_repeat_scrapes(client, druks_db, monkeypatch) -> None: + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) fetched: list[str] = [] monkeypatch.setattr(ClaudeHarness, "fetch_usage", _fake_fetch(fetched)) diff --git a/backend/tests/test_app_doctor_checks.py b/backend/tests/test_app_doctor_checks.py index e70d9121..bda61252 100644 --- a/backend/tests/test_app_doctor_checks.py +++ b/backend/tests/test_app_doctor_checks.py @@ -1,3 +1,4 @@ +from contextlib import asynccontextmanager from pathlib import Path import pytest @@ -7,6 +8,14 @@ from druks.testing import make_settings from druks.user_settings.models import SettingsOverride + +@asynccontextmanager +async def _fixture_check_engine(_settings): + from druks.database import db_session + + yield db_session().bind + + # field_notes is the out-of-tree proof app (``backend/tests/druks-field_notes``). # It declares settings coherence and one check on its class. These tests drive both # through the platform's own doctor without doctor importing the app's private @@ -24,7 +33,7 @@ def _named(results: list[doctor.CheckResult], name: str) -> doctor.CheckResult: return next(result for result in results if result.name == name) -def test_passing_app_check_reports_under_the_app( +async def test_passing_app_check_reports_under_the_app( installed, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """A satisfied precondition passes and its result is namespaced under the app @@ -32,44 +41,46 @@ def test_passing_app_check_reports_under_the_app( monkeypatch.setenv("FIELD_NOTES_API_KEY", "sk-test") settings = make_settings(tmp_path) - result = _named(doctor.check_apps(settings), "field_notes:summary_api_key") + result = _named(await doctor.check_apps(settings), "field_notes:summary_api_key") assert result.ok assert result.detail == "set" -def test_failing_app_check_reports_under_the_app(installed, tmp_path: Path) -> None: +async def test_failing_app_check_reports_under_the_app(installed, tmp_path: Path) -> None: """The app's API-key check fails when the credential is unset, reported under the app name so the operator knows which app is broken.""" settings = make_settings(tmp_path) - result = _named(doctor.check_apps(settings), "field_notes:summary_api_key") + result = _named(await doctor.check_apps(settings), "field_notes:summary_api_key") assert not result.ok assert "FIELD_NOTES_API_KEY" in result.detail -def test_unreachable_settings_database_is_an_app_check_failure(installed, tmp_path: Path) -> None: +async def test_unreachable_settings_database_is_an_app_check_failure( + installed, tmp_path: Path +) -> None: settings = make_settings( tmp_path, database_url="postgresql+psycopg://druks:druks@127.0.0.1:1/druks", ) - results = doctor.check_apps(settings) + results = await doctor.check_apps(settings) result = _named(results, "ship:settings") assert not result.ok assert "check raised" in result.detail -def test_selected_unconnected_tracker_pends_through_ships_own_check( +async def test_selected_unconnected_tracker_pends_through_ships_own_check( installed, tmp_path: Path, druks_db, monkeypatch: pytest.MonkeyPatch ) -> None: # The default selector names linear; no identity is connected in this db. - monkeypatch.setattr(doctor, "create_engine_from_url", lambda _: druks_db.get_bind()) + monkeypatch.setattr(doctor, "_check_engine", _fixture_check_engine) try: - result = _named(doctor.check_apps(make_settings(tmp_path)), "ship:tracker") + result = _named(await doctor.check_apps(make_settings(tmp_path)), "ship:tracker") finally: db_session.registry.set(druks_db) @@ -78,17 +89,17 @@ def test_selected_unconnected_tracker_pends_through_ships_own_check( assert "linear" in result.detail -def test_half_configured_review_identity_fails_through_review_settings( +async def test_half_configured_review_identity_fails_through_review_settings( installed, tmp_path: Path, druks_db, monkeypatch: pytest.MonkeyPatch ) -> None: # Review identity health is Review's own: the incoherent pair fails under # ``review:settings`` while the set/unset check stays healthy — no core # doctor check hardcodes review knowledge, and no GitHub call is made. - SettingsOverride.set_app_setting("review", "app_id", "42", is_secret=True) - monkeypatch.setattr(doctor, "create_engine_from_url", lambda _: druks_db.get_bind()) + await SettingsOverride.set_app_setting("review", "app_id", "42", is_secret=True) + monkeypatch.setattr(doctor, "_check_engine", _fixture_check_engine) try: - results = doctor.check_apps(make_settings(tmp_path)) + results = await doctor.check_apps(make_settings(tmp_path)) finally: db_session.registry.set(druks_db) @@ -100,16 +111,16 @@ def test_half_configured_review_identity_fails_through_review_settings( assert _named(results, "review:identity").ok -def test_coherent_stored_settings_pass( +async def test_coherent_stored_settings_pass( installed, tmp_path: Path, druks_db, monkeypatch: pytest.MonkeyPatch ) -> None: - SettingsOverride.set_app_setting( + await SettingsOverride.set_app_setting( "ship", "linear_trigger_status", "Agent Queue", is_secret=False ) - monkeypatch.setattr(doctor, "create_engine_from_url", lambda _: druks_db.get_bind()) + monkeypatch.setattr(doctor, "_check_engine", _fixture_check_engine) try: - result = _named(doctor.check_apps(make_settings(tmp_path)), "ship:settings") + result = _named(await doctor.check_apps(make_settings(tmp_path)), "ship:settings") finally: db_session.registry.set(druks_db) @@ -117,7 +128,7 @@ def test_coherent_stored_settings_pass( assert result.detail == "coherent" -def test_app_without_settings_has_no_settings_row( +async def test_app_without_settings_has_no_settings_row( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: class Plain(App): @@ -125,10 +136,10 @@ class Plain(App): monkeypatch.setattr(doctor, "iter_apps", lambda: iter([Plain])) - assert doctor.check_apps(make_settings(tmp_path)) == [] + assert await doctor.check_apps(make_settings(tmp_path)) == [] -def test_raising_settings_clean_is_contained( +async def test_raising_settings_clean_is_contained( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: class Broken(App): @@ -140,13 +151,13 @@ def clean(self) -> dict[str, str]: monkeypatch.setattr(doctor, "iter_apps", lambda: iter([Broken])) - result = _named(doctor.check_apps(make_settings(tmp_path)), "broken_settings:settings") + result = _named(await doctor.check_apps(make_settings(tmp_path)), "broken_settings:settings") assert not result.ok assert "coherence crashed" in result.detail -def test_app_checks_are_wired_into_the_check_battery(installed, tmp_path: Path) -> None: +async def test_app_checks_are_wired_into_the_check_battery(installed, tmp_path: Path) -> None: """``run_checks`` runs the app checks: ``check_apps`` is one of the battery's entries and, like ``check_harness_credentials``, fans its several results into the run — so the app's checks reach the report beside core's.""" @@ -154,12 +165,12 @@ def test_app_checks_are_wired_into_the_check_battery(installed, tmp_path: Path) assert doctor.check_apps in doctor.CHECKS - app_results = doctor.check_apps(settings) + app_results = await doctor.check_apps(settings) assert isinstance(app_results, list) assert "field_notes:summary_api_key" in {result.name for result in app_results} -def test_raising_app_check_is_isolated_and_does_not_stop_siblings( +async def test_raising_app_check_is_isolated_and_does_not_stop_siblings( installed, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """A check that raises becomes one failing result tagged with the app name, @@ -175,7 +186,7 @@ def healthy() -> doctor.CheckResult: monkeypatch.setattr(field_notes.FieldNotes, "checks", [boom, healthy]) settings = make_settings(tmp_path) - results = doctor.check_apps(settings) + results = await doctor.check_apps(settings) raised = _named(results, "field_notes:boom") assert not raised.ok @@ -184,7 +195,7 @@ def healthy() -> doctor.CheckResult: assert _named(results, "field_notes:healthy").ok -def test_broken_app_check_does_not_hide_core_failures( +async def test_broken_app_check_does_not_hide_core_failures( installed, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """The key robustness contract: a raising app check is contained inside @@ -205,7 +216,7 @@ def boom() -> doctor.CheckResult: assert doctor.check_redis in doctor.CHECKS # The app's raising check is contained as a failure under its own name… - app_result = _named(doctor.check_apps(settings), "field_notes:boom") + app_result = _named(await doctor.check_apps(settings), "field_notes:boom") assert not app_result.ok assert "kaboom" in app_result.detail @@ -226,7 +237,7 @@ class Plain(App): assert Plain.checks == [] -def test_malformed_check_return_is_contained( +async def test_malformed_check_return_is_contained( installed, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """A check that returns something other than a ``CheckResult`` — a missing @@ -244,7 +255,7 @@ def healthy() -> doctor.CheckResult: monkeypatch.setattr(field_notes.FieldNotes, "checks", [check_forgot_return, healthy]) settings = make_settings(tmp_path) - results = doctor.check_apps(settings) + results = await doctor.check_apps(settings) by_name = {result.name: result for result in results} # The malformed return is contained as a failure under its own name… diff --git a/backend/tests/test_apps.py b/backend/tests/test_apps.py index 19977f16..038eb237 100644 --- a/backend/tests/test_apps.py +++ b/backend/tests/test_apps.py @@ -14,7 +14,7 @@ class Widget(Subject): in a real app; these stand in for that.""" @classmethod - def list_summaries(cls, account_id: str | None) -> list: + async def list_summaries(cls, account_id: str | None) -> list: return [] diff --git a/backend/tests/test_artifacts.py b/backend/tests/test_artifacts.py index 7974dfb9..c16f0a63 100644 --- a/backend/tests/test_artifacts.py +++ b/backend/tests/test_artifacts.py @@ -10,26 +10,26 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert -def _seed_call(druks_db) -> AgentCall: +async def _seed_call(druks_db) -> AgentCall: druks_db.add(Run(id="run-1", kind="build")) call = AgentCall( id="call-1", run_id="run-1", agent="summarize", model="m", sandbox_host_id="host-1" ) druks_db.add(call) - druks_db.flush() + await druks_db.flush() return call -def test_record_writes_content_and_descriptor(druks_db, tmp_path): - _seed_call(druks_db) - Artifact.record( +async def test_record_writes_content_and_descriptor(druks_db, tmp_path): + await _seed_call(druks_db) + await Artifact.record( call_dir=tmp_path, call_id="call-1", kind="markdown", title="Implementation plan", content="# Plan\nbody", ) - artifact = Artifact.get_for_call("call-1") + artifact = await Artifact.get_for_call("call-1") assert artifact is not None assert (artifact.kind, artifact.title, artifact.path) == ( "markdown", @@ -39,28 +39,32 @@ def test_record_writes_content_and_descriptor(druks_db, tmp_path): assert (tmp_path / "artifact.md").read_text() == "# Plan\nbody" -def test_record_is_idempotent_per_call(druks_db, tmp_path): +async def test_record_is_idempotent_per_call(druks_db, tmp_path): # A replayed step must not double-record; the unique fk makes record a no-op. - _seed_call(druks_db) + await _seed_call(druks_db) for _ in range(2): - Artifact.record( + await Artifact.record( call_dir=tmp_path, call_id="call-1", kind="markdown", title="P", content="x" ) - rows = druks_db.scalars(select(Artifact).where(Artifact.agent_call_id == "call-1")).all() + rows = ( + await druks_db.scalars(select(Artifact).where(Artifact.agent_call_id == "call-1")) + ).all() assert len(rows) == 1 -def test_artifact_cascades_with_its_call(druks_db, tmp_path): - call = _seed_call(druks_db) - Artifact.record(call_dir=tmp_path, call_id="call-1", kind="markdown", title="P", content="x") - assert Artifact.get_for_call("call-1") is not None +async def test_artifact_cascades_with_its_call(druks_db, tmp_path): + call = await _seed_call(druks_db) + await Artifact.record( + call_dir=tmp_path, call_id="call-1", kind="markdown", title="P", content="x" + ) + assert await Artifact.get_for_call("call-1") is not None - druks_db.delete(call) - druks_db.flush() - assert Artifact.get_for_call("call-1") is None + await druks_db.delete(call) + await druks_db.flush() + assert await Artifact.get_for_call("call-1") is None -def test_get_latest_for_run_returns_the_newest_calls_artifact(druks_db, tmp_path): +async def test_get_latest_for_run_returns_the_newest_calls_artifact(druks_db, tmp_path): # The read side serves the run's latest artifact on the in-app review ask — # the second call's plan wins. druks_db.add(Run(id="run-1", kind="build")) @@ -70,19 +74,19 @@ def test_get_latest_for_run_returns_the_newest_calls_artifact(druks_db, tmp_path id=call_id, run_id="run-1", agent="summarize", model="m", sandbox_host_id="host-1" ) ) - druks_db.flush() - Artifact.record( + await druks_db.flush() + await Artifact.record( call_dir=tmp_path / call_id, call_id=call_id, kind="markdown", title=title, content="x", ) - latest = Artifact.get_latest_for_run("run-1") + latest = await Artifact.get_latest_for_run("run-1") assert latest is not None and latest.title == "Revised plan" -def test_get_ask_resolves_the_review_artifact(druks_db, tmp_path): +async def test_get_ask_resolves_the_review_artifact(druks_db, tmp_path): # An in-app ask stores no label/artifact — the read side derives both from # the run's latest artifact. A declared ask passes through untouched. run = Run( @@ -97,15 +101,19 @@ def test_get_ask_resolves_the_review_artifact(druks_db, tmp_path): id="call-1", run_id="run-1", agent="summarize", model="m", sandbox_host_id="host-1" ) ) - druks_db.flush() - Artifact.record(call_dir=tmp_path, call_id="call-1", kind="markdown", title="Plan", content="x") + await druks_db.flush() + await Artifact.record( + call_dir=tmp_path, call_id="call-1", kind="markdown", title="Plan", content="x" + ) - ask = RunResponse.from_run(run, input_request=run.get_ask()).input_request + await druks_db.refresh(run) + await run.awaitable_attrs.agent_calls + ask = RunResponse.from_run(run, input_request=await run.get_ask()).input_request assert ask == { "presentation": "in_app", "controls": ["approve"], "label": "Review: Plan", - "artifact_id": Artifact.get_for_call("call-1").id, + "artifact_id": (await Artifact.get_for_call("call-1")).id, } external = Run( @@ -115,17 +123,20 @@ def test_get_ask_resolves_the_review_artifact(druks_db, tmp_path): input_request={"presentation": "external", "label": "Review implementation"}, ) druks_db.add(external) - druks_db.flush() - response = RunResponse.from_run(external, input_request=external.get_ask()) + await druks_db.flush() + await druks_db.refresh(external) + await external.awaitable_attrs.agent_calls + response = RunResponse.from_run(external, input_request=await external.get_ask()) assert response.input_request == { "presentation": "external", "label": "Review implementation", } -def test_run_response_projects_the_parked_gate(druks_db): - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-gate", input_gate="review") +async def test_run_response_projects_the_parked_gate(druks_db): + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-gate", input_gate="review") + await run.awaitable_attrs.agent_calls response = RunResponse.from_run(run, input_request=None) assert response.gate == "review" @@ -138,15 +149,15 @@ async def test_get_artifact_returns_recorded_content(druks_db, tmp_path, monkeyp id="call-1", run_id="run-1", agent="summarize", model="m", sandbox_host_id="host-1" ) druks_db.add(call) - druks_db.flush() - Artifact.record( + await druks_db.flush() + await Artifact.record( call_dir=call.call_dir, call_id="call-1", kind="markdown", title="Implementation plan", content="# Plan\nbody", ) - result = await get_artifact(Artifact.get_for_call("call-1").id) + result = await get_artifact((await Artifact.get_for_call("call-1")).id) assert (result.kind, result.title, result.content) == ( "markdown", "Implementation plan", @@ -170,13 +181,13 @@ async def test_get_artifact_404_when_content_gone(druks_db, tmp_path, monkeypatc id="call-1", run_id="run-1", agent="summarize", model="m", sandbox_host_id="host-1" ) ) - druks_db.flush() - druks_db.execute( + await druks_db.flush() + await druks_db.execute( pg_insert(Artifact).values( id="art-1", agent_call_id="call-1", kind="markdown", title="P", path="artifact.md" ) ) - druks_db.flush() + await druks_db.flush() with pytest.raises(HTTPException) as exc: await get_artifact("art-1") assert exc.value.status_code == 404 diff --git a/backend/tests/test_attribution.py b/backend/tests/test_attribution.py index 101bddfe..eca5e681 100644 --- a/backend/tests/test_attribution.py +++ b/backend/tests/test_attribution.py @@ -5,22 +5,24 @@ from druks_field_notes.workflows import Summarize -def test_run_projects_its_account(druks_db): - account = Account.get_or_create("dev@example.com") - seed_run(druks_db, kind=Summarize.kind, run_id="run-attr-1", account_id=account.id) - druks_db.flush() +async def test_run_projects_its_account(druks_db): + account = await Account.get_or_create("dev@example.com") + await seed_run(druks_db, kind=Summarize.kind, run_id="run-attr-1", account_id=account.id) + await druks_db.flush() - run = druks_db.get(Run, "run-attr-1") + run = await druks_db.get(Run, "run-attr-1") + await run.awaitable_attrs.agent_calls assert run.account_id == account.id response = RunResponse.from_run(run, input_request=None) assert response.account_username == "dev@example.com" -def test_an_unowned_run_belongs_to_system(druks_db): - seed_run(druks_db, kind=Summarize.kind, run_id="run-attr-2") - druks_db.flush() +async def test_an_unowned_run_belongs_to_system(druks_db): + await seed_run(druks_db, kind=Summarize.kind, run_id="run-attr-2") + await druks_db.flush() - run = druks_db.get(Run, "run-attr-2") + run = await druks_db.get(Run, "run-attr-2") + await run.awaitable_attrs.agent_calls assert run.account_id == "system" response = RunResponse.from_run(run, input_request=None) assert response.account_username == "system" diff --git a/backend/tests/test_auth_pats.py b/backend/tests/test_auth_pats.py index 5ce534d4..7e5dd4ec 100644 --- a/backend/tests/test_auth_pats.py +++ b/backend/tests/test_auth_pats.py @@ -29,13 +29,13 @@ def _bearer(token: str) -> dict[str, str]: return {"Authorization": f"Bearer {token}"} -def _mint(username: str = "agent@example.com") -> tuple[PersonalAccessToken, str]: - account = Account.get_or_create(username) - return PersonalAccessToken.create(account_id=account.id, name="agent") +async def _mint(username: str = "agent@example.com") -> tuple[PersonalAccessToken, str]: + account = await Account.get_or_create(username) + return await PersonalAccessToken.create(account_id=account.id, name="agent") -def test_the_minted_token_shape_and_hash_are_pinned(druks_db): - pat, token = _mint() +async def test_the_minted_token_shape_and_hash_are_pinned(druks_db): + pat, token = await _mint() prefix, _, secret = token.removeprefix(f"{PAT_TOKEN_TAG}_").partition("_") assert token.startswith(f"{PAT_TOKEN_TAG}_") assert len(prefix) == 12 @@ -49,8 +49,8 @@ def test_the_minted_token_shape_and_hash_are_pinned(druks_db): assert pat.status == "active" -def test_a_prefix_collision_regenerates(druks_db, monkeypatch): - first, _ = _mint() +async def test_a_prefix_collision_regenerates(druks_db, monkeypatch): + first, _ = await _mint() replay = iter(first.token_prefix) random_choice = secrets.choice @@ -62,43 +62,43 @@ def collide_once(alphabet): return random_choice(alphabet) monkeypatch.setattr(secrets, "choice", collide_once) - second, _ = PersonalAccessToken.create(account_id=first.account_id, name="two") + second, _ = await PersonalAccessToken.create(account_id=first.account_id, name="two") assert second.token_prefix != first.token_prefix -def test_authenticate_rejects_everything_but_the_live_token(druks_db): - pat, token = _mint() - assert PersonalAccessToken.authenticate(token).id == pat.id +async def test_authenticate_rejects_everything_but_the_live_token(druks_db): + pat, token = await _mint() + assert (await PersonalAccessToken.authenticate(token)).id == pat.id with pytest.raises(InvalidPatError): - PersonalAccessToken.authenticate("not-even-shaped-right") + await PersonalAccessToken.authenticate("not-even-shaped-right") with pytest.raises(InvalidPatError): - PersonalAccessToken.authenticate(f"{PAT_TOKEN_TAG}_{pat.token_prefix}_wrongsecret") + await PersonalAccessToken.authenticate(f"{PAT_TOKEN_TAG}_{pat.token_prefix}_wrongsecret") pat.expires_at = Base.utc_now() - timedelta(days=1) with pytest.raises(InvalidPatError, match=f"{pat.token_prefix} has expired"): - PersonalAccessToken.authenticate(token) + await PersonalAccessToken.authenticate(token) pat.expires_at = Base.utc_now() + timedelta(days=1) - pat.revoke() + await pat.revoke() with pytest.raises(InvalidPatError, match=f"{pat.token_prefix} was revoked"): - PersonalAccessToken.authenticate(token) + await PersonalAccessToken.authenticate(token) -def test_last_used_advances_at_most_hourly(druks_db): - pat, token = _mint() - PersonalAccessToken.authenticate(token) +async def test_last_used_advances_at_most_hourly(druks_db): + pat, token = await _mint() + await PersonalAccessToken.authenticate(token) first_use = pat.last_used_at assert first_use - PersonalAccessToken.authenticate(token) + await PersonalAccessToken.authenticate(token) assert pat.last_used_at == first_use pat.last_used_at = first_use - timedelta(hours=2) - PersonalAccessToken.authenticate(token) + await PersonalAccessToken.authenticate(token) assert pat.last_used_at > first_use - timedelta(hours=2) -def test_a_bearer_pat_authenticates_gated_routes(tmp_path, druks_db): +async def test_a_bearer_pat_authenticates_gated_routes(tmp_path, druks_db): with _client(tmp_path) as client: - _, token = _mint() + _, token = await _mint() response = client.get("/api/auth/me", headers=_bearer(token)) assert response.status_code == 200 assert response.json()["account"]["username"] == "agent@example.com" @@ -126,7 +126,7 @@ def test_a_shapeless_authorization_header_is_challenged(tmp_path, druks_db, head @pytest.mark.parametrize("header", ["Bearer a b", "bearer lowercased", "BEARER nope"]) -def test_any_scheme_case_reaches_authentication_and_fails_closed(tmp_path, druks_db, header): +async def test_any_scheme_case_reaches_authentication_and_fails_closed(tmp_path, druks_db, header): # RFC 7235 schemes are case-insensitive: these parse as credentials and die # in authentication — a 401 either way, never a slide to the assertion. with _client(tmp_path) as client: @@ -135,7 +135,7 @@ def test_any_scheme_case_reaches_authentication_and_fails_closed(tmp_path, druks ) assert response.status_code == 401 assert response.headers["WWW-Authenticate"].endswith('error="invalid_token"') - assert not Account.get_for_username("op@example.com") + assert not await Account.get_for_username("op@example.com") def test_an_empty_authorization_header_never_slides_to_the_assertion(tmp_path, druks_db): @@ -144,10 +144,10 @@ def test_an_empty_authorization_header_never_slides_to_the_assertion(tmp_path, d assert response.status_code == 401 -def test_a_dead_token_401s_with_its_prefix_only(tmp_path, druks_db): +async def test_a_dead_token_401s_with_its_prefix_only(tmp_path, druks_db): with _client(tmp_path) as client: - pat, token = _mint() - pat.revoke() + pat, token = await _mint() + await pat.revoke() response = client.get("/api/auth/me", headers=_bearer(token)) assert response.status_code == 401 assert response.headers["WWW-Authenticate"] == 'Bearer realm="druks", error="invalid_token"' @@ -156,9 +156,9 @@ def test_a_dead_token_401s_with_its_prefix_only(tmp_path, druks_db): assert secret not in response.text -def test_a_pat_cannot_manage_pats(tmp_path, druks_db): +async def test_a_pat_cannot_manage_pats(tmp_path, druks_db): with _client(tmp_path) as client: - pat, token = _mint() + pat, token = await _mint() assert client.get("/api/auth/personal-tokens", headers=_bearer(token)).status_code == 401 create = client.post( "/api/auth/personal-tokens", json={"name": "x"}, headers=_bearer(token) @@ -172,11 +172,11 @@ def test_a_pat_cannot_manage_pats(tmp_path, druks_db): assert both.status_code == 401 -def test_a_pat_cannot_disconnect_a_harness(tmp_path, druks_db): +async def test_a_pat_cannot_disconnect_a_harness(tmp_path, druks_db): # Disconnect destroys a capability a bearer could never create — the same # session-only rule as token management. with _client(tmp_path) as client: - _, token = _mint("op@example.com") + _, token = await _mint("op@example.com") alone = client.delete("/api/harnesses/claude/connection", headers=_bearer(token)) assert alone.status_code == 401 beside = client.delete( @@ -185,9 +185,9 @@ def test_a_pat_cannot_disconnect_a_harness(tmp_path, druks_db): assert beside.status_code == 401 -def test_a_pat_reads_but_cannot_write_app_settings(tmp_path, druks_db): +async def test_a_pat_reads_but_cannot_write_app_settings(tmp_path, druks_db): with _client(tmp_path) as client: - _, token = _mint("op@example.com") + _, token = await _mint("op@example.com") headers = _bearer(token) read = client.get("/api/settings/apps", headers=headers) @@ -201,9 +201,9 @@ def test_a_pat_reads_but_cannot_write_app_settings(tmp_path, druks_db): assert write.status_code == 401 -def test_a_pat_reads_but_cannot_write_service_identities(tmp_path, druks_db): +async def test_a_pat_reads_but_cannot_write_service_identities(tmp_path, druks_db): with _client(tmp_path) as client: - _, token = _mint("op@example.com") + _, token = await _mint("op@example.com") headers = _bearer(token) read = client.get("/api/services", headers=headers) @@ -242,8 +242,8 @@ def test_the_operator_manages_the_token_lifecycle(tmp_path, druks_db): assert again["revokedAt"] == revoked["revokedAt"] -def test_the_none_mode_operator_manages_tokens_too(tmp_path, druks_db): - Account.get_or_create("op@example.com") +async def test_the_none_mode_operator_manages_tokens_too(tmp_path, druks_db): + await Account.get_or_create("op@example.com") with _client(tmp_path, identity={"mode": "none"}) as client: created = client.post("/api/auth/personal-tokens", json={"name": "local"}) assert created.status_code == 200 @@ -251,20 +251,20 @@ def test_the_none_mode_operator_manages_tokens_too(tmp_path, druks_db): assert [item["name"] for item in listed] == ["local"] -def test_the_list_is_scoped_to_the_operator(tmp_path, druks_db): +async def test_the_list_is_scoped_to_the_operator(tmp_path, druks_db): with _client(tmp_path) as client: - _mint("other@example.com") + await _mint("other@example.com") assert client.get("/api/auth/personal-tokens", headers=OPERATOR).json() == [] -def test_revoking_anothers_token_is_a_404(tmp_path, druks_db): +async def test_revoking_anothers_token_is_a_404(tmp_path, druks_db): with _client(tmp_path) as client: - pat, _ = _mint("other@example.com") + pat, _ = await _mint("other@example.com") assert ( client.delete(f"/api/auth/personal-tokens/{pat.id}", headers=OPERATOR).status_code == 404 ) - session_registry().expire_all() + session_registry().expunge_all() assert not pat.revoked_at diff --git a/backend/tests/test_browser_borrow.py b/backend/tests/test_browser_borrow.py index f1a8d47b..544b7666 100644 --- a/backend/tests/test_browser_borrow.py +++ b/backend/tests/test_browser_borrow.py @@ -105,15 +105,15 @@ async def ephemeral(*, image_override, provider): return browser, redis -def stored_session( +async def stored_session( declaration: BrowserSession, payload: bytes = b"stored-state" ) -> StoredBrowserSession: - row = StoredBrowserSession.get_or_create( + row = await StoredBrowserSession.get_or_create( name=declaration.name, payload_format=BrowserSessionPayloadFormat.STORAGE_STATE, site=declaration.site, ) - row.store_payload(payload) + await row.store_payload(payload) return row @@ -124,7 +124,7 @@ def test_declaration_carries_the_app_namespace(night_watch): async def test_borrow_yields_a_tunneled_cdp_url(borrow, night_watch): browser, redis = borrow - stored_session(night_watch.docs) + await stored_session(night_watch.docs) async with night_watch.docs.cdp() as cdp_url: assert cdp_url == "http://127.0.0.1:43987" @@ -140,7 +140,7 @@ async def test_borrow_yields_a_tunneled_cdp_url(borrow, night_watch): launch_script = browser.commands[0][2] assert "session-launch --headed" in launch_script assert not redis.values - assert StoredBrowserSession.get_for_name(night_watch.docs.name).last_used_at + assert (await StoredBrowserSession.get_for_name(night_watch.docs.name)).last_used_at async def test_headless_declaration_launches_headless(borrow): @@ -148,7 +148,7 @@ async def test_headless_declaration_launches_headless(borrow): quiet = BrowserSession(site="docs.example") quiet.headless = True quiet.name = "night_watch.quiet" - stored_session(quiet) + await stored_session(quiet) async with quiet.cdp(): pass @@ -158,15 +158,15 @@ async def test_headless_declaration_launches_headless(borrow): async def test_persisting_borrow_locks_exports_and_stores(borrow, night_watch): browser, redis = borrow - row = stored_session(night_watch.acme) + row = await stored_session(night_watch.acme) async with night_watch.acme.cdp(): assert redis.values assert not redis.values assert browser.commands[-1] == ["session-export"] - db_session().expire_all() - stored = StoredBrowserSession.get_for_name(night_watch.acme.name) + db_session().expunge_all() + stored = await StoredBrowserSession.get_for_name(night_watch.acme.name) assert stored.payload.decrypt() == b"exported-profile" assert stored.payload_format == BrowserSessionPayloadFormat.PROFILE_DIR.value assert stored.id == row.id @@ -174,9 +174,9 @@ async def test_persisting_borrow_locks_exports_and_stores(borrow, night_watch): async def test_persisting_borrow_refuses_a_second_writer(borrow, night_watch): browser, redis = borrow - stored_session(night_watch.acme) + await stored_session(night_watch.acme) redis.values[ - f"browser_session:{StoredBrowserSession.get_for_name(night_watch.acme.name).id}" + f"browser_session:{(await StoredBrowserSession.get_for_name(night_watch.acme.name)).id}" ] = "other" with pytest.raises(BrowserSessionWriterLockedError): @@ -191,13 +191,13 @@ async def test_first_borrow_writes_the_declared_session_and_asks_for_a_login( ): """The first borrow materializes the row and refuses to open a browser: the session is declared, but nobody has signed into it yet.""" - assert not StoredBrowserSession.get_for_name(night_watch.docs.name) + assert not await StoredBrowserSession.get_for_name(night_watch.docs.name) with pytest.raises(BrowserSessionNotReadyError): async with night_watch.docs.cdp(): pass - row = StoredBrowserSession.get_for_name(night_watch.docs.name) + row = await StoredBrowserSession.get_for_name(night_watch.docs.name) assert row.status == BrowserSessionStatus.NEEDS_LOGIN.value assert row.site == night_watch.docs.site @@ -205,12 +205,12 @@ async def test_first_borrow_writes_the_declared_session_and_asks_for_a_login( async with night_watch.docs.cdp(): pass - assert StoredBrowserSession.list_all() == [row] + assert await StoredBrowserSession.list_all() == [row] async def test_launch_failure_raises_and_releases_the_lock(borrow, night_watch): browser, redis = borrow - stored_session(night_watch.acme) + await stored_session(night_watch.acme) browser.launch_exit = 1 with pytest.raises(BrowserLaunchError, match="launch stderr"): @@ -224,17 +224,17 @@ async def test_signed_out_borrow_stamps_the_session_and_stores_nothing(borrow, n """The app raises through the borrow when the site bounced the login: the door stamps which session bounced, and the dead state is never stored.""" browser, redis = borrow - stored_session(night_watch.acme, payload=b"live-state") + await stored_session(night_watch.acme, payload=b"live-state") with pytest.raises(BrowserSessionSignedOutError) as caught: async with night_watch.acme.cdp(): raise BrowserSessionSignedOutError("the site bounced the login") assert caught.value.session_name == "night_watch.acme" - db_session().expire_all() + db_session().expunge_all() assert ( - StoredBrowserSession.get_for_name(night_watch.acme.name).payload.decrypt() == b"live-state" - ) + await StoredBrowserSession.get_for_name(night_watch.acme.name) + ).payload.decrypt() == b"live-state" assert ["session-export"] not in browser.commands assert not redis.values # the writer lock released on the way out @@ -254,7 +254,7 @@ async def test_anonymous_borrow_needs_no_login(borrow, night_watch): } assert ["session-export"] not in browser.commands assert not redis.values # no writer lock: nothing to serialize - row = StoredBrowserSession.get_for_name(night_watch.status_page.name) + row = await StoredBrowserSession.get_for_name(night_watch.status_page.name) assert row.status == BrowserSessionStatus.ANONYMOUS.value assert row.last_used_at assert not row.payload @@ -274,7 +274,7 @@ async def test_signed_out_in_an_anonymous_borrow_keeps_the_row_anonymous(borrow, assert caught.value.session_name == "night_watch.status_page" await signed_out_session_goes_stale(session_name="night_watch.status_page") - row = StoredBrowserSession.get_for_name("night_watch.status_page") + row = await StoredBrowserSession.get_for_name("night_watch.status_page") assert row.status == BrowserSessionStatus.ANONYMOUS.value @@ -284,7 +284,7 @@ async def test_playwright_yields_the_logged_in_context(borrow, night_watch, monk from contextlib import asynccontextmanager as acm browser, _ = borrow - stored_session(night_watch.docs) + await stored_session(night_watch.docs) seen = {} logged_in_context = object() @@ -317,7 +317,7 @@ async def fake_playwright(): async def test_playwright_without_the_dependency_names_the_fix(borrow, night_watch, monkeypatch): import sys - stored_session(night_watch.docs) + await stored_session(night_watch.docs) monkeypatch.setitem(sys.modules, "playwright", None) monkeypatch.setitem(sys.modules, "playwright.async_api", None) diff --git a/backend/tests/test_browser_session_login_window.py b/backend/tests/test_browser_session_login_window.py index 86b1f82a..ce5a3786 100644 --- a/backend/tests/test_browser_session_login_window.py +++ b/backend/tests/test_browser_session_login_window.py @@ -82,8 +82,8 @@ def window_runtime(tmp_path, monkeypatch): return client -def create_session(name: str = "acme-main") -> StoredBrowserSession: - return StoredBrowserSession.get_or_create( +async def create_session(name: str = "acme-main") -> StoredBrowserSession: + return await StoredBrowserSession.get_or_create( name=name, payload_format=BrowserSessionPayloadFormat.STORAGE_STATE, site="acme.example", @@ -93,7 +93,7 @@ def create_session(name: str = "acme-main") -> StoredBrowserSession: async def test_login_launch_leaves_the_box_untouched_when_nothing_is_set(window_runtime): client = window_runtime - await LoginWindow.open(create_session()) + await LoginWindow.open(await create_session()) command = client.browsers[0].launch_command or "" assert "DRUKS_BROWSER_LOGIN_PROXY" not in command @@ -103,7 +103,7 @@ async def test_login_launch_leaves_the_box_untouched_when_nothing_is_set(window_ async def test_login_launch_opens_on_the_session_site(window_runtime): client = window_runtime - await LoginWindow.open(create_session()) + await LoginWindow.open(await create_session()) command = client.browsers[0].launch_command or "" assert "DRUKS_BROWSER_URL=https://acme.example" in command @@ -122,7 +122,7 @@ async def test_login_launch_routes_through_the_configured_proxy(tmp_path, monkey proxy = "http://172.17.0.1:8888" client = _runtime_with_sandbox(tmp_path, monkeypatch, browser_login_proxy=proxy) - await LoginWindow.open(create_session()) + await LoginWindow.open(await create_session()) command = client.browsers[0].launch_command or "" assert f"DRUKS_BROWSER_LOGIN_PROXY={shlex.quote(proxy)}" in command @@ -131,7 +131,7 @@ async def test_login_launch_routes_through_the_configured_proxy(tmp_path, monkey async def test_login_launch_sets_the_configured_timezone(tmp_path, monkeypatch): client = _runtime_with_sandbox(tmp_path, monkeypatch, browser_login_tz="Europe/Madrid") - await LoginWindow.open(create_session()) + await LoginWindow.open(await create_session()) command = client.browsers[0].launch_command or "" assert "TZ=Europe/Madrid" in command @@ -141,7 +141,7 @@ async def test_login_launch_quotes_values_so_a_bad_one_cannot_inject(tmp_path, m proxy = "http://h:8888; rm -rf /" client = _runtime_with_sandbox(tmp_path, monkeypatch, browser_login_proxy=proxy) - await LoginWindow.open(create_session()) + await LoginWindow.open(await create_session()) command = client.browsers[0].launch_command or "" assert shlex.quote(proxy) in command @@ -151,7 +151,7 @@ async def test_login_launch_quotes_values_so_a_bad_one_cannot_inject(tmp_path, m async def test_open_seeds_a_blank_profile_and_records_the_container(window_runtime): client = window_runtime - session = create_session() + session = await create_session() await LoginWindow.open(session) @@ -170,7 +170,7 @@ async def test_open_seeds_a_blank_profile_and_records_the_container(window_runti async def test_reopening_disposes_the_previous_window(window_runtime): client = window_runtime - session = create_session() + session = await create_session() await LoginWindow.open(session) await LoginWindow.open(session) @@ -181,9 +181,9 @@ async def test_reopening_disposes_the_previous_window(window_runtime): async def test_storage_state_reconnect_saves_a_profile(window_runtime): client = window_runtime - session = create_session() - session.store_payload(b'{"cookies": [{"name": "login"}], "origins": []}') - session.mark_stale() + session = await create_session() + await session.store_payload(b'{"cookies": [{"name": "login"}], "origins": []}') + await session.mark_stale() await LoginWindow.open(session) browser = client.browsers[0] @@ -196,8 +196,8 @@ async def test_storage_state_reconnect_saves_a_profile(window_runtime): saved = await (await LoginWindow.get_for_session(session.name)).save() assert saved.payload_format == BrowserSessionPayloadFormat.PROFILE_DIR - db_session().expire_all() - stored = StoredBrowserSession.get_for_name(session.name) + db_session().expunge_all() + stored = await StoredBrowserSession.get_for_name(session.name) assert stored.status == BrowserSessionStatus.READY.value assert stored.payload.decrypt() == b"fresh-profile" assert client.released == [browser.id] @@ -207,7 +207,7 @@ async def test_storage_state_reconnect_saves_a_profile(window_runtime): async def test_failed_export_closes_the_window(window_runtime): client = window_runtime - session = create_session() + session = await create_session() await LoginWindow.open(session) client.browsers[0].export_exit_code = 1 @@ -221,7 +221,7 @@ async def test_failed_export_closes_the_window(window_runtime): async def test_cancel_then_cancel_again_reports_the_window_gone(window_runtime): client = window_runtime - session = create_session() + session = await create_session() await LoginWindow.open(session) await (await LoginWindow.get_for_session(session.name)).cancel() diff --git a/backend/tests/test_browser_sessions.py b/backend/tests/test_browser_sessions.py index e0027b05..07452862 100644 --- a/backend/tests/test_browser_sessions.py +++ b/backend/tests/test_browser_sessions.py @@ -46,7 +46,9 @@ async def open(cls, session) -> None: cls.opened.append(session.name) -def test_declared_sessions_list_without_a_row_and_the_pane_read_writes_nothing(client, night_watch): +async def test_declared_sessions_list_without_a_row_and_the_pane_read_writes_nothing( + client, night_watch +): listed = client.get("/api/browser-sessions").json() assert [entry["name"] for entry in listed] == ["night_watch.acme", "night_watch.docs"] @@ -56,11 +58,11 @@ def test_declared_sessions_list_without_a_row_and_the_pane_read_writes_nothing(c assert entry["payloadFormat"] is None assert entry["createdAt"] is None assert entry["site"] == "acme.example" - assert not StoredBrowserSession.list_all() + assert not await StoredBrowserSession.list_all() -def test_leftover_rows_list_as_undeclared_and_refuse_the_login_window(client, night_watch): - StoredBrowserSession.get_or_create( +async def test_leftover_rows_list_as_undeclared_and_refuse_the_login_window(client, night_watch): + await StoredBrowserSession.get_or_create( name="gone_ext.old", payload_format=BrowserSessionPayloadFormat.PROFILE_DIR, site="gone.example", @@ -75,10 +77,10 @@ def test_leftover_rows_list_as_undeclared_and_refuse_the_login_window(client, ni assert client.post("/api/browser-sessions/gone_ext.old/login-window").status_code == 404 assert client.delete("/api/browser-sessions/gone_ext.old").status_code == 204 - assert not StoredBrowserSession.list_all() + assert not await StoredBrowserSession.list_all() -def test_anonymous_sessions_list_as_anonymous_and_refuse_login_and_state( +async def test_anonymous_sessions_list_as_anonymous_and_refuse_login_and_state( client, browser_session_declarations ): class Critic: @@ -96,10 +98,12 @@ class Critic: "/api/browser-sessions/critic.target/state?payloadFormat=storage_state", content=b"x" ) assert uploaded.status_code == 409 - assert not StoredBrowserSession.list_all() + assert not await StoredBrowserSession.list_all() -def test_opening_the_login_window_materializes_the_declared_row(client, night_watch, monkeypatch): +async def test_opening_the_login_window_materializes_the_declared_row( + client, night_watch, monkeypatch +): monkeypatch.setattr(routes, "LoginWindow", FakeLoginWindow) monkeypatch.setattr(FakeLoginWindow, "opened", []) @@ -107,14 +111,14 @@ def test_opening_the_login_window_materializes_the_declared_row(client, night_wa assert opened.status_code == 204 assert FakeLoginWindow.opened == ["night_watch.acme"] - row = StoredBrowserSession.get_for_name("night_watch.acme") + row = await StoredBrowserSession.get_for_name("night_watch.acme") assert row.status == BrowserSessionStatus.NEEDS_LOGIN.value assert row.site == "acme.example" assert client.post("/api/browser-sessions/nobody.home/login-window").status_code == 404 -def test_import_materializes_the_row_survives_restart_and_delete_removes_it( +async def test_import_materializes_the_row_survives_restart_and_delete_removes_it( client, night_watch, tmp_path, monkeypatch ): payload = b'{"cookies":[{"name":"auth_token","value":"secret"}],"origins":[]}' @@ -129,15 +133,13 @@ def test_import_materializes_the_row_survives_restart_and_delete_removes_it( assert listed["night_watch.acme"]["payloadFormat"] == BrowserSessionPayloadFormat.STORAGE_STATE assert listed["night_watch.acme"]["lastRefreshedAt"] - row = StoredBrowserSession.get_for_name("night_watch.acme") + row = await StoredBrowserSession.get_for_name("night_watch.acme") stored = ( - db_session() - .execute( + await db_session().execute( text("SELECT payload FROM browser_sessions WHERE id = :id"), {"id": row.id}, ) - .scalar_one() - ) + ).scalar_one() assert payload not in bytes(stored) with pytest.raises(SecretDecryptError): secret_utils.decrypt(bytes(stored), "another_table.payload") @@ -150,8 +152,8 @@ def test_import_materializes_the_row_survives_restart_and_delete_removes_it( with pytest.raises(SecretDecryptError): row.payload.decrypt() - db_session().expire_all() - restarted = StoredBrowserSession.get_for_name("night_watch.acme") + db_session().expunge_all() + restarted = await StoredBrowserSession.get_for_name("night_watch.acme") assert restarted.payload.decrypt() == payload undeclared = client.put( @@ -161,7 +163,7 @@ def test_import_materializes_the_row_survives_restart_and_delete_removes_it( deleted = client.delete("/api/browser-sessions/night_watch.acme") assert deleted.status_code == 204 - assert not StoredBrowserSession.list_all() + assert not await StoredBrowserSession.list_all() def test_upload_rejects_payloads_above_the_cap(client, night_watch, monkeypatch): @@ -190,14 +192,14 @@ def test_upload_warns_at_the_product_threshold(client, night_watch, monkeypatch, assert "received a 5-byte payload" in caplog.text -def test_bearer_pat_reads_sessions_but_cannot_mutate_them(tmp_path, druks_db, night_watch): +async def test_bearer_pat_reads_sessions_but_cannot_mutate_them(tmp_path, druks_db, night_watch): settings = make_settings( tmp_path, identity={"mode": "header", "header": "X-Edge-Email"}, ) - account = Account.get_or_create("op@example.com") - _, token = PersonalAccessToken.create(account_id=account.id, name="agent") - db_session().commit() + account = await Account.get_or_create("op@example.com") + _, token = await PersonalAccessToken.create(account_id=account.id, name="agent") + await db_session().commit() headers = {"Authorization": f"Bearer {token}"} with TestClient(configure_app_for_test(settings=settings, authenticated=False)) as pat_client: diff --git a/backend/tests/test_cost_capture.py b/backend/tests/test_cost_capture.py index e087273a..97c13fbb 100644 --- a/backend/tests/test_cost_capture.py +++ b/backend/tests/test_cost_capture.py @@ -76,29 +76,31 @@ def test_read_cost_returns_none_for_corrupt_file(tmp_path: Path): assert metadata is None -def test_record_agent_run_cost_persists_to_db(druks_db): - note = Note.create(body="cost capture") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", status="running") - call.record_cost( +async def test_record_agent_run_cost_persists_to_db(druks_db): + note = await Note.create(body="cost capture") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", status="running") + # record_cost flushes the ambient session, so mutate the ambient copy. + call = await AgentCall.get(call.id) + await call.record_cost( cost_usd=1.23, cost_metadata={"input_tokens": 200, "model": "claude-opus-4-7"}, ) - fetched = AgentCall.get(call.id) + fetched = await AgentCall.get(call.id) assert fetched is not None assert fetched.cost_usd == 1.23 assert fetched.cost_metadata == {"input_tokens": 200, "model": "claude-opus-4-7"} -def test_record_agent_run_cost_noop_when_empty(druks_db): - note = Note.create(body="empty cost") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", status="running") +async def test_record_agent_run_cost_noop_when_empty(druks_db): + note = await Note.create(body="empty cost") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", status="running") - call.record_cost(cost_usd=None, cost_metadata=None) + await call.record_cost(cost_usd=None, cost_metadata=None) - fetched = AgentCall.get(call.id) + fetched = await AgentCall.get(call.id) assert fetched is not None assert fetched.cost_usd is None assert fetched.cost_metadata is None diff --git a/backend/tests/test_doctor.py b/backend/tests/test_doctor.py index c5b4a9b4..a6df8a51 100644 --- a/backend/tests/test_doctor.py +++ b/backend/tests/test_doctor.py @@ -1,3 +1,4 @@ +from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta from pathlib import Path @@ -8,6 +9,14 @@ from druks.services.models import ServiceIdentity from druks.testing import make_settings + +@asynccontextmanager +async def _fixture_check_engine(_settings): + from druks.database import db_session + + yield db_session().bind + + _SECRETS_KEY = "MDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDA=" @@ -20,13 +29,13 @@ def doctor_db(druks_db, monkeypatch: pytest.MonkeyPatch): """Point doctor's one-off engines at the test transaction. Doctor's checks remove their session from the ambient registry, so rebind the fixture's afterwards for the teardown that still needs it.""" - monkeypatch.setattr(doctor, "create_engine_from_url", lambda _url: druks_db.get_bind()) + monkeypatch.setattr(doctor, "_check_engine", _fixture_check_engine) yield druks_db db_session.registry.set(druks_db) -def _connect_github(slug: str = "druks-operator") -> ServiceIdentity: - return ServiceIdentity.connect( +async def _connect_github(slug: str = "druks-operator") -> ServiceIdentity: + return await ServiceIdentity.connect( "github", identity={"app_id": "12345", "slug": slug}, secrets={ @@ -40,30 +49,32 @@ def _github_identity_result(results: list[doctor.CheckResult]) -> doctor.CheckRe return next(result for result in results if result.name == "github_identity") -def test_service_identities_pending_when_a_required_service_is_absent( +async def test_service_identities_pending_when_a_required_service_is_absent( tmp_path: Path, doctor_db ) -> None: - result = _github_identity_result(doctor.check_service_identities(make_settings(tmp_path))) + result = _github_identity_result(await doctor.check_service_identities(make_settings(tmp_path))) assert not result.ok assert result.pending assert "not connected" in result.detail -def test_service_identities_report_the_connected_row(tmp_path: Path, doctor_db) -> None: - _connect_github() +async def test_service_identities_report_the_connected_row(tmp_path: Path, doctor_db) -> None: + await _connect_github() - result = _github_identity_result(doctor.check_service_identities(make_settings(tmp_path))) + result = _github_identity_result(await doctor.check_service_identities(make_settings(tmp_path))) assert result.ok assert "app_id=12345" in result.detail assert "slug=druks-operator" in result.detail -def test_installations_pending_without_a_connected_identity(tmp_path: Path, doctor_db) -> None: +async def test_installations_pending_without_a_connected_identity( + tmp_path: Path, doctor_db +) -> None: # No github row → the zero-argument client can't even be built; doctor # reports it as pending operator setup instead of raising. - result = doctor.check_installations(make_settings(tmp_path)) + result = await doctor.check_installations(make_settings(tmp_path)) assert not result.ok assert result.pending @@ -71,39 +82,47 @@ def test_installations_pending_without_a_connected_identity(tmp_path: Path, doct assert "not connected" in result.detail -def test_installations_lists_accounts(tmp_path: Path, doctor_db, monkeypatch) -> None: +async def test_installations_lists_accounts(tmp_path: Path, doctor_db, monkeypatch) -> None: class _FakeClient: async def list_installation_accounts(self): return ("clawhaven",) - monkeypatch.setattr("druks.doctor.get_github_client", lambda: _FakeClient()) + async def _fake_client(): + return _FakeClient() - result = doctor.check_installations(make_settings(tmp_path)) + monkeypatch.setattr("druks.doctor.get_github_client", _fake_client) + + result = await doctor.check_installations(make_settings(tmp_path)) assert result.ok assert "clawhaven" in result.detail -def test_installations_pending_when_app_has_none(tmp_path: Path, doctor_db, monkeypatch) -> None: +async def test_installations_pending_when_app_has_none( + tmp_path: Path, doctor_db, monkeypatch +) -> None: class _FakeClient: async def list_installation_accounts(self): return () - monkeypatch.setattr("druks.doctor.get_github_client", lambda: _FakeClient()) + async def _fake_client(): + return _FakeClient() + + monkeypatch.setattr("druks.doctor.get_github_client", _fake_client) - result = doctor.check_installations(make_settings(tmp_path)) + result = await doctor.check_installations(make_settings(tmp_path)) assert not result.ok assert result.pending assert "no installations" in result.detail -def test_installations_builds_the_client_from_the_row( +async def test_installations_builds_the_client_from_the_row( tmp_path: Path, doctor_db, monkeypatch ) -> None: # The real zero-argument factory resolves the row inside doctor's own # bound session; the fake transport keeps GitHub out of it. - _connect_github() + await _connect_github() class _FakeClient: async def list_installation_accounts(self): @@ -112,14 +131,14 @@ async def list_installation_accounts(self): real_factory = doctor.get_github_client built: list[str] = [] - def _tracking_factory(): - client = real_factory() + async def _tracking_factory(): + client = await real_factory() built.append(client._app_id) return _FakeClient() monkeypatch.setattr(doctor, "get_github_client", _tracking_factory) - result = doctor.check_installations(make_settings(tmp_path)) + result = await doctor.check_installations(make_settings(tmp_path)) assert result.ok assert built == ["12345"] @@ -164,21 +183,21 @@ def test_redis_fails_on_unreachable_host(tmp_path: Path) -> None: assert "127.0.0.1:1" in result.detail -def test_drukbox_passes_when_unconfigured(tmp_path: Path) -> None: +async def test_drukbox_passes_when_unconfigured(tmp_path: Path) -> None: """Sandbox URL empty → no drukbox to talk to.""" settings = make_settings(tmp_path) assert settings.sandbox.service_url == "" - result = doctor.check_drukbox(settings) + result = await doctor.check_drukbox(settings) assert result.ok assert "not configured" in result.detail -def test_run_checks_covers_all_check_names(tmp_path: Path) -> None: +async def test_run_checks_covers_all_check_names(tmp_path: Path) -> None: settings = make_settings(tmp_path) - results = doctor.run_checks(settings) + results = await doctor.run_checks(settings) # Installed apps contribute their own checks alongside the platform's. assert {result.name for result in results} >= { @@ -346,44 +365,46 @@ async def release(self, *, host_id): return calls -def test_sandbox_e2e_not_configured_is_ok(tmp_path: Path) -> None: +async def test_sandbox_e2e_not_configured_is_ok(tmp_path: Path) -> None: settings = make_settings(tmp_path, sandbox={"service_url": ""}) - result = doctor.check_sandbox_e2e(settings) + result = await doctor.check_sandbox_e2e(settings) assert result.ok assert result.detail == "not configured" -def test_sandbox_e2e_exercises_dial_and_reattach(tmp_path: Path, monkeypatch) -> None: +async def test_sandbox_e2e_exercises_dial_and_reattach(tmp_path: Path, monkeypatch) -> None: calls = _fake_sandbox_client(monkeypatch) settings = make_settings(tmp_path, sandbox={"service_url": "http://127.0.0.1:8780"}) - result = doctor.check_sandbox_e2e(settings) + result = await doctor.check_sandbox_e2e(settings) assert result.ok assert "reattach" in result.detail assert calls == ["acquire", "attach:host-doc", "release:host-doc"] -def test_sandbox_e2e_failure_names_the_phase_and_releases(tmp_path: Path, monkeypatch) -> None: +async def test_sandbox_e2e_failure_names_the_phase_and_releases( + tmp_path: Path, monkeypatch +) -> None: """A reattach failure is the bug class worth this check — the error surfaces in the detail, and the VM must still be released.""" calls = _fake_sandbox_client(monkeypatch, reattach_fails=True) settings = make_settings(tmp_path, sandbox={"service_url": "http://127.0.0.1:8780"}) - result = doctor.check_sandbox_e2e(settings) + result = await doctor.check_sandbox_e2e(settings) assert not result.ok assert "dial timed out" in result.detail assert "release:host-doc" in calls -def test_run_checks_includes_sandbox_e2e_only_when_flagged(tmp_path: Path) -> None: +async def test_run_checks_includes_sandbox_e2e_only_when_flagged(tmp_path: Path) -> None: settings = make_settings(tmp_path, sandbox={"service_url": ""}) - default = {r.name for r in doctor.run_checks(settings)} - flagged = {r.name for r in doctor.run_checks(settings, sandbox=True)} + default = {r.name for r in await doctor.run_checks(settings)} + flagged = {r.name for r in await doctor.run_checks(settings, sandbox=True)} assert "sandbox_e2e" not in default assert "sandbox_e2e" in flagged diff --git a/backend/tests/test_durable_schemas.py b/backend/tests/test_durable_schemas.py index 4470ea63..4cba12a2 100644 --- a/backend/tests/test_durable_schemas.py +++ b/backend/tests/test_durable_schemas.py @@ -27,55 +27,55 @@ def _run( ) -def _status_of(runs): +async def _status_of(runs): # runs arrives newest-first, mirroring Run.list_for_subject. - return _status(runs[0]) + return await _status(runs[0]) -def test_subject_state_takes_the_newest_run(): +async def test_subject_state_takes_the_newest_run(): runs = [ _run("new", "ship.build", RunState.RUNNING), _run("old", "ship.build", RunState.PARKED), ] - assert _status_of(runs).state == RunState.RUNNING + assert (await _status_of(runs)).state == RunState.RUNNING -def test_subject_state_takes_a_newer_parked_run_over_an_older_running_one(): +async def test_subject_state_takes_a_newer_parked_run_over_an_older_running_one(): # Recency decides, not a hardcoded state preference. runs = [ _run("new", "ship.build", RunState.PARKED), _run("old", "ship.build", RunState.RUNNING), ] - assert _status_of(runs).state == RunState.PARKED + assert (await _status_of(runs)).state == RunState.PARKED -def test_subject_state_uses_the_latest_outcome_once_every_run_is_terminal(): +async def test_subject_state_uses_the_latest_outcome_once_every_run_is_terminal(): runs = [ _run("new", "ship.build", RunState.FINISHED), _run("old", "ship.build", RunState.FAILED), ] - assert _status_of(runs).state == RunState.FINISHED + assert (await _status_of(runs)).state == RunState.FINISHED -def test_status_surfaces_the_newest_active_runs_gate(): +async def test_status_surfaces_the_newest_active_runs_gate(): runs = [ _run("new", "ship.build", RunState.PARKED, "review"), _run("old", "ship.build", RunState.PARKED, "review_work"), ] - assert _status_of(runs).gate == "review" + assert (await _status_of(runs)).gate == "review" -def test_status_carries_the_running_runs_kind_and_no_stale_gate(): +async def test_status_carries_the_running_runs_kind_and_no_stale_gate(): runs = [ _run("new", "ship.build", RunState.RUNNING), _run("old", "ship.build", RunState.PARKED, "review"), ] - status = _status_of(runs) + status = await _status_of(runs) assert status.kind == "ship.build" assert not status.gate -def test_status_carries_the_latest_agent_call_agent(): +async def test_status_carries_the_latest_agent_call_agent(): runs = [ _run( "new", @@ -84,10 +84,10 @@ def test_status_carries_the_latest_agent_call_agent(): agent_calls=[AgentCall(agent="generate_plan"), AgentCall(agent="implement")], ) ] - assert _status_of(runs).agent == "implement" + assert (await _status_of(runs)).agent == "implement" -def test_parked_status_carries_no_agent_even_when_the_run_has_calls(): +async def test_parked_status_carries_no_agent_even_when_the_run_has_calls(): # A parked run keeps its calls, but the status never reads them — the fact # stays consistent with the board, where a parked row never queries them. runs = [ @@ -99,32 +99,32 @@ def test_parked_status_carries_no_agent_even_when_the_run_has_calls(): agent_calls=[AgentCall(agent="implement")], ) ] - status = _status_of(runs) + status = await _status_of(runs) assert not status.agent assert status.gate == "review" -def test_status_carries_the_gate_timeout_reason(): +async def test_status_carries_the_gate_timeout_reason(): # The gate timeout's stamped failure_code rides the status as ``reason`` — # the board renders the re-trigger hint from it instead of a bare "failed". runs = [ _run("new", Summarize.kind, RunState.FAILED, failure_code=GateTimeout.code), ] - status = _status_of(runs) + status = await _status_of(runs) assert status.reason == GateTimeout.code assert status.kind == Summarize.kind -def test_status_carries_failure_but_no_reason_when_the_run_crashed(): +async def test_status_carries_failure_but_no_reason_when_the_run_crashed(): runs = [ _run("new", Summarize.kind, RunState.FAILED, failure="boom"), ] - status = _status_of(runs) + status = await _status_of(runs) assert not status.reason assert status.failure == "boom" -def test_status_falls_back_to_the_terminal_call_error_when_failure_is_null(): +async def test_status_falls_back_to_the_terminal_call_error_when_failure_is_null(): # A crash leaves run.failure null; the terminal call's captured error becomes # the status reason so the board never shows a bare "failed". runs = [ @@ -135,4 +135,4 @@ def test_status_falls_back_to_the_terminal_call_error_when_failure_is_null(): agent_calls=[AgentCall(agent="implement", last_error="crashed in implement")], ) ] - assert _status_of(runs).failure == "crashed in implement" + assert (await _status_of(runs)).failure == "crashed in implement" diff --git a/backend/tests/test_durable_sdk.py b/backend/tests/test_durable_sdk.py index 920ef09a..5f414177 100644 --- a/backend/tests/test_durable_sdk.py +++ b/backend/tests/test_durable_sdk.py @@ -15,7 +15,8 @@ from druks.testing import init_db from druks.workflows import Gate, Subject, Workflow, step, task from pydantic import BaseModel -from sqlalchemy import create_engine, select +from sqlalchemy import NullPool, create_engine, select +from sqlalchemy.ext.asyncio import create_async_engine PG_BASE = os.environ.get("DRUKS_TEST_PG", "postgresql://druks:druks@localhost:5432") DB = "druks_durable_test" @@ -175,7 +176,7 @@ class ScheduledDispatch(Workflow): @classmethod async def dispatch(cls) -> str: - return await cls.start(subject=Widget.get_for_subject_id("313131")) + return await cls.start(subject=await Widget.get_for_subject_id("313131")) async def run(self) -> None: ... @@ -185,7 +186,7 @@ class SubjectFlow(Workflow): subject = Widget async def run(self) -> Decision: - SINK.append(f"subj-id:{self.subject.id}") + SINK.append(f"subj-id:{(await self.subject).id}") return Decision(action="ok") class DoubleGateFlow(Workflow): @@ -247,7 +248,7 @@ async def run_multistep(self) -> None: @pytest.fixture(scope="module", autouse=True) -def rt(): +async def rt(): db_url_snap = os.environ.get("DRUKS_DATABASE_URL") admin = psycopg.connect(f"{PG_BASE}/postgres", autocommit=True) @@ -255,8 +256,9 @@ def rt(): admin.execute(f"CREATE DATABASE {DB}") admin.close() - engine = create_engine(URL) - init_db(engine) # full schema incl. durable_runs + the work_items chain + schema_engine = create_engine(URL) + init_db(schema_engine) # full schema incl. durable_runs + the work_items chain + engine = create_async_engine(URL, poolclass=NullPool) configure_engine(engine) configure_session(engine) @@ -272,7 +274,7 @@ def rt(): try: account = Account(username="op@example.com") session.add(account) - session.flush() + await session.flush() session.add_all( Widget(id=subject_id) for subject_id in (7, 4242, 636363, 424242, 515151, 878787, 909090, 313131) @@ -285,10 +287,12 @@ def rt(): payload={"claudeAiOauth": {"accessToken": "t"}}, ) ) - session.merge(UserSettings(id=UserSettings.SINGLETON_ID, fallback_account_id=account.id)) - session.commit() + await session.merge( + UserSettings(id=UserSettings.SINGLETON_ID, fallback_account_id=account.id) + ) + await session.commit() finally: - session.close() + await session.close() ( sample_flow, @@ -310,10 +314,11 @@ def rt(): ) = _build_units() os.environ["DRUKS_DATABASE_URL"] = URL init_dbos() - launch() # also runs apply_schedules() for daily_sweep + await launch() # also runs await apply_schedules() for daily_sweep try: yield SimpleNamespace( engine=engine, + schema_engine=schema_engine, SampleFlow=sample_flow, AgentFlow=agent_flow, AgentBodyFlow=agent_body_flow, @@ -333,7 +338,8 @@ def rt(): ) finally: shutdown() - engine.dispose() + await engine.dispose() + schema_engine.dispose() # Drop only the test's own keys so other modules see clean registries # (a wholesale restore would clobber registrations made meanwhile). agents._items.pop("decider", None) @@ -357,37 +363,38 @@ def rt(): os.environ["DRUKS_DATABASE_URL"] = db_url_snap -def _state(engine, workflow_id: str) -> Run | None: +async def _state(engine, workflow_id: str) -> Run | None: session = get_session(engine) try: - return session.get(Run, workflow_id) + return await session.get(Run, workflow_id) finally: - session.close() + await session.close() async def _wait_for(engine, workflow_id, predicate, timeout=15.0): deadline = asyncio.get_event_loop().time() + timeout while asyncio.get_event_loop().time() < deadline: - row = _state(engine, workflow_id) + row = await _state(engine, workflow_id) if row is not None and predicate(row): return row await asyncio.sleep(0.1) - raise AssertionError(f"timed out; last={_state(engine, workflow_id)}") + raise AssertionError(f"timed out; last={await _state(engine, workflow_id)}") -def _account_id(engine, email: str) -> str: +async def _account_id(engine, email: str) -> str: from druks.accounts.models import Account session = get_session(engine) try: - row = session.execute(select(Account).where(Account.username == email)).scalar_one_or_none() + result = await session.execute(select(Account).where(Account.username == email)) + row = result.scalar_one_or_none() if not row: row = Account(username=email) session.add(row) - session.commit() + await session.commit() return row.id finally: - session.close() + await session.close() async def test_attribution_rides_the_run_and_survives_resume(rt): @@ -396,10 +403,10 @@ async def test_attribution_rides_the_run_and_survives_resume(rt): from druks.durable.dbos_state import workflow_status SINK.clear() - account_id = _account_id(rt.engine, "op@example.com") + account_id = await _account_id(rt.engine, "op@example.com") wfid = await rt.AttributedFlow.start(subject=Widget(id=878787), account_id=account_id) parked = await _wait_for(rt.engine, wfid, lambda r: r.state == RunState.PARKED) - with rt.engine.connect() as conn: + with rt.schema_engine.connect() as conn: attributes = conn.execute( select(workflow_status.c.attributes).where(workflow_status.c.workflow_uuid == wfid) ).scalar_one() @@ -421,14 +428,14 @@ async def test_browser_origin_start_inherits_the_ambient_account(rt): # start() reads it when no explicit account_id is passed. from druks.accounts.context import current_account_id - account_id = _account_id(rt.engine, "ambient@example.com") + account_id = await _account_id(rt.engine, "ambient@example.com") token = current_account_id.set(account_id) try: wfid = await rt.RecordFeedback.start(subject=None, repo="owner/ambient") finally: current_account_id.reset(token) await _wait_for(rt.engine, wfid, lambda r: r.state == RunState.FINISHED) - assert _state(rt.engine, wfid).account_id == account_id + assert (await _state(rt.engine, wfid)).account_id == account_id async def test_duplicate_start_shares_the_run_across_accounts(rt): @@ -444,8 +451,8 @@ async def test_duplicate_start_shares_the_run_across_accounts(rt): async def saw(*, subject=None, **_: object) -> None: scheduled.append(subject) - first = _account_id(rt.engine, "op@example.com") - second = _account_id(rt.engine, "peer@example.com") + first = await _account_id(rt.engine, "op@example.com") + second = await _account_id(rt.engine, "peer@example.com") subject = Widget(id=909090) wfid = await rt.SampleFlow.start(subject=subject, account_id=first, repo="owner/app") parked = await _wait_for(rt.engine, wfid, lambda r: r.state == RunState.PARKED) @@ -487,7 +494,7 @@ async def test_duplicate_replies_to_one_round_collapse(rt): await parked.resume(action="duplicate") # The duplicate collapsed against the round's one notification. - with rt.engine.connect() as conn: + with rt.schema_engine.connect() as conn: delivered = conn.execute( text( "SELECT count(*) FROM dbos.notifications" @@ -530,7 +537,7 @@ async def test_fail_branch(rt): assert failed.failure == "closed at review" # FAILED derives from DBOS's own record: the FatalError re-raised out of the # workflow, so DBOS wrote terminal ERROR, not SUCCESS. - with rt.engine.connect() as conn: + with rt.schema_engine.connect() as conn: status = conn.execute( text("SELECT status FROM dbos.workflow_status WHERE workflow_uuid = :id"), {"id": wfid}, @@ -555,9 +562,9 @@ async def test_signed_out_run_fails_and_marks_the_session_stale(rt): site="acme.example", ) ) - session.commit() + await session.commit() finally: - session.close() + await session.close() class BounceFlow(Workflow): async def run(self) -> None: @@ -572,12 +579,16 @@ async def run(self) -> None: assert failed.failure_code == "browser_session_signed_out" session = get_session(rt.engine) try: - stored = session.execute( - select(StoredBrowserSession).where(StoredBrowserSession.name == "night_watch.acme") + stored = ( + await session.execute( + select(StoredBrowserSession).where( + StoredBrowserSession.name == "night_watch.acme" + ) + ) ).scalar_one() assert stored.status == BrowserSessionStatus.STALE.value finally: - session.close() + await session.close() finally: workflows._items.pop("bounce_flow", None) @@ -601,7 +612,7 @@ async def test_subject_gate_parks_unchanged(rt): assert parked.input_gate == "confirm" # start() stamped the subject as workflow attributes — the keying every # runs-for-a-subject query reads; the id normalizes to a string. - with rt.engine.connect() as conn: + with rt.schema_engine.connect() as conn: attributes = conn.execute( select(workflow_status.c.attributes).where(workflow_status.c.workflow_uuid == wfid) ).scalar_one() @@ -688,14 +699,16 @@ async def test_run_agent_step(rt, monkeypatch): assert seen[0]["agent"] == "decider" session = get_session(rt.engine) try: - recorded = list(session.query(AgentCall).filter(AgentCall.run_id == wfid)) + recorded = list( + (await session.execute(select(AgentCall).where(AgentCall.run_id == wfid))).scalars() + ) finally: - session.close() + await session.close() # The call is recorded under the orchestrator-minted id threaded to run_agent. assert recorded[0].id == seen[0]["call_id"] # No account on the start: the fallback account (the module's op@ seed) # is charged. - assert recorded[0].account_id == _account_id(rt.engine, "op@example.com") + assert recorded[0].account_id == await _account_id(rt.engine, "op@example.com") assert held == [False] # the step let its connection go before the agent ran @@ -828,18 +841,18 @@ async def test_scheduled_tick_fires_dispatch_not_run(rt): _, fn = next(row for row in _scheduled if row[0].kind == "scheduled_dispatch") await fn(datetime.now(UTC), None) - def dispatched_run(): + async def dispatched_run(): session = get_session(rt.engine) try: - return session.execute( - select(Run).where(Run.kind == "scheduled_dispatch") + return ( + await session.execute(select(Run).where(Run.kind == "scheduled_dispatch")) ).scalar_one_or_none() finally: - session.close() + await session.close() deadline = asyncio.get_event_loop().time() + 15 while asyncio.get_event_loop().time() < deadline: - run = dispatched_run() + run = await dispatched_run() if run and run.state == RunState.FINISHED: break await asyncio.sleep(0.1) @@ -874,7 +887,7 @@ async def test_apply_schedules_drops_undeclared(rt): DBOS.create_schedule(schedule_name="stale_cron", workflow_fn=fn, schedule=cls.every) assert "stale_cron" in {s["schedule_name"] for s in DBOS.list_schedules()} - apply_schedules() + await apply_schedules() live = {s["schedule_name"] for s in DBOS.list_schedules()} assert "stale_cron" not in live # undeclared → dropped @@ -897,20 +910,20 @@ def sweep_cron(): # Each write commits — a bare test-task session stays idle-in-transaction # and its row locks deadlock any later test touching the same rows. - with session_scope(rt.engine): - SettingsOverride.write("workflow:daily_sweep:schedule", "0 9 * * *") - apply_schedules() + async with session_scope(rt.engine): + await SettingsOverride.write("workflow:daily_sweep:schedule", "0 9 * * *") + await apply_schedules() assert sweep_cron() == "0 9 * * *" # override wins over the declared default - with session_scope(rt.engine): - SettingsOverride.write("workflow:daily_sweep:schedule_enabled", False) - apply_schedules() + async with session_scope(rt.engine): + await SettingsOverride.write("workflow:daily_sweep:schedule_enabled", False) + await apply_schedules() assert sweep_cron() is None # paused → no schedule, nothing fires - with session_scope(rt.engine): - SettingsOverride.write("workflow:daily_sweep:schedule", None) - SettingsOverride.write("workflow:daily_sweep:schedule_enabled", None) - apply_schedules() + async with session_scope(rt.engine): + await SettingsOverride.write("workflow:daily_sweep:schedule", None) + await SettingsOverride.write("workflow:daily_sweep:schedule_enabled", None) + await apply_schedules() assert sweep_cron() == "0 6 * * *" # overrides cleared → declared default @@ -920,16 +933,16 @@ async def test_session_scope_commits_writes(rt): from druks.database import session_scope from druks.user_settings.models import SettingsOverride - with session_scope(rt.engine): - SettingsOverride.write("session_scope_commit_probe", {"landed": True}) + async with session_scope(rt.engine): + await SettingsOverride.write("session_scope_commit_probe", {"landed": True}) session = get_session(rt.engine) try: - row = session.get(SettingsOverride, "session_scope_commit_probe") + row = await session.get(SettingsOverride, "session_scope_commit_probe") assert row is not None assert row.value == {"landed": True} finally: - session.close() + await session.close() async def test_launch_commits_the_user_settings_seed(rt): @@ -941,9 +954,9 @@ async def test_launch_commits_the_user_settings_seed(rt): session = get_session(rt.engine) try: - assert session.get(UserSettings, UserSettings.SINGLETON_ID) is not None + assert await session.get(UserSettings, UserSettings.SINGLETON_ID) is not None finally: - session.close() + await session.close() async def test_apply_schedules_evaluates_cron_in_operator_timezone(rt): @@ -957,16 +970,16 @@ def sweep_timezone(): rows = {s["schedule_name"]: s["cron_timezone"] for s in DBOS.list_schedules()} return rows.get("daily_sweep") - apply_schedules() + await apply_schedules() assert sweep_timezone() == "UTC" # the settings default # Commit the write — a bare test-task session stays idle-in-transaction and # its row lock deadlocks any later test that touches user_settings. from druks.database import session_scope - with session_scope(rt.engine): - UserSettings.get().update_profile(timezone="Europe/Madrid") - apply_schedules() + async with session_scope(rt.engine): + await (await UserSettings.get()).update_profile(timezone="Europe/Madrid") + await apply_schedules() assert sweep_timezone() == "Europe/Madrid" @@ -977,27 +990,27 @@ async def test_user_settings_get_recreates_the_singleton(rt): from druks.user_settings.models import UserSettings from sqlalchemy import delete - with session_scope(rt.engine): - db_session().execute(delete(UserSettings)) - with session_scope(rt.engine): - assert UserSettings.get().timezone == "UTC" - with session_scope(rt.engine): - assert UserSettings.get().id == UserSettings.SINGLETON_ID + async with session_scope(rt.engine): + await db_session().execute(delete(UserSettings)) + async with session_scope(rt.engine): + assert (await UserSettings.get()).timezone == "UTC" + async with session_scope(rt.engine): + assert (await UserSettings.get()).id == UserSettings.SINGLETON_ID async def test_a_run_hydrates_the_subject_row_it_was_started_for(rt): from druks.database import db_session, session_scope - with session_scope(rt.engine): + async with session_scope(rt.engine): widget = Widget() db_session().add(widget) - db_session().flush() + await db_session().flush() assert widget.identity == {"type": "widget", "id": widget.id} run = rt.SubjectFlow() run._subject = widget.identity - assert run.subject is widget + assert await run.subject is widget async def test_input_is_validated_at_start(rt): @@ -1095,12 +1108,12 @@ async def test_subject_reaches_body_and_result_rides_finished_event(rt): session = get_session(rt.engine) try: finished = ( - session.query(Event) - .filter(Event.type == "workflow.finished", Event.subject_id == "7") - .one() - ) + await session.execute( + select(Event).where(Event.type == "workflow.finished", Event.subject_id == "7") + ) + ).scalar_one() finally: - session.close() + await session.close() # run()'s BaseModel return rides the finished event. assert finished.payload["result"] == {"action": "ok"} @@ -1128,9 +1141,15 @@ async def test_run_events_carry_subject(rt): session = get_session(rt.engine) try: - events = list(session.query(Event).filter(Event.subject_id == "4242").order_by(Event.id)) + events = list( + ( + await session.execute( + select(Event).where(Event.subject_id == "4242").order_by(Event.id) + ) + ).scalars() + ) finally: - session.close() + await session.close() assert [e.type for e in events] == ["workflow.running", "workflow.finished"] assert {e.subject_type for e in events} == {"widget"} @@ -1191,8 +1210,9 @@ async def test_subjectless_run_emits_no_events(rt): session = get_session(rt.engine) try: - events = [e for e in session.query(Event).all() if e.payload.get("run") == wfid] + rows = (await session.execute(select(Event))).scalars() + events = [e for e in rows if e.payload.get("run") == wfid] finally: - session.close() + await session.close() assert events == [] diff --git a/backend/tests/test_events_feed.py b/backend/tests/test_events_feed.py index b7d56041..69a2bd01 100644 --- a/backend/tests/test_events_feed.py +++ b/backend/tests/test_events_feed.py @@ -16,19 +16,19 @@ class Pallet(StoredSubject): __tablename__ = "faketest_pallets" -def test_feed_carries_what_a_row_is_worded_from(druks_db): - note = Note.create(body="the pump ran hot") - Event.emit( +async def test_feed_carries_what_a_row_is_worded_from(druks_db): + note = await Note.create(body="the pump ran hot") + await Event.emit( type="workflow.running", subject=note.identity, label=note.label, app="field_notes", payload={"kind": Summarize.kind, "run": "wf1"}, ) - Event.emit(type="summarized", subject=note.identity, label=note.label, app="field_notes") - druks_db.flush() + await Event.emit(type="summarized", subject=note.identity, label=note.label, app="field_notes") + await druks_db.flush() - by_kind = {row.kind: row for row in build_feed()[0]} + by_kind = {row.kind: row for row in (await build_feed())[0]} started = by_kind["workflow.running"] assert (started.app, started.workflow) == ("field_notes", Summarize.kind) @@ -40,36 +40,38 @@ def test_feed_carries_what_a_row_is_worded_from(druks_db): assert by_kind["summarized"].workflow is None -def test_every_subject_shows_itself(druks_db): +async def test_every_subject_shows_itself(druks_db): # A subject that declares a handle reads as it; one that doesn't reads by # identity. Either way it is snapshotted, so the row survives the row itself. crate, pallet = Crate(id=7), Pallet(id=7) druks_db.add_all([crate, pallet]) - druks_db.flush() + await druks_db.flush() assert crate.identity == {"type": "crate", "id": 7} for subject in (crate, pallet): - Event.emit(type="stocked", subject=subject.identity, label=subject.label, app="faketest") - druks_db.delete(crate) - druks_db.flush() + await Event.emit( + type="stocked", subject=subject.identity, label=subject.label, app="faketest" + ) + await druks_db.delete(crate) + await druks_db.flush() - by_type = {row.subject_type: row for row in build_feed()[0] if row.app == "faketest"} + by_type = {row.subject_type: row for row in (await build_feed())[0] if row.app == "faketest"} assert by_type["crate"].subject_label == "CRATE-7" assert by_type["pallet"].subject_label == "pallet 7" -def test_feed_paginates_same_second_events_without_loss_or_repeat(druks_db): +async def test_feed_paginates_same_second_events_without_loss_or_repeat(druks_db): # utc_now truncates to whole seconds, so these all share a created_at. Paging on # the truncated timestamp used to drop the whole second on the next page; paging on # the monotonic pk covers every event exactly once. for i in range(5): - Event.emit(type=f"evt-{i}") - druks_db.flush() + await Event.emit(type=f"evt-{i}") + await druks_db.flush() collected = [] cursor = None for _ in range(10): # bounded so a paging bug can't loop forever - page, cursor = build_feed(before=int(cursor) if cursor else None, limit=2) + page, cursor = await build_feed(before=int(cursor) if cursor else None, limit=2) collected.extend(page) if cursor is None: break diff --git a/backend/tests/test_gate.py b/backend/tests/test_gate.py index 92976fb2..7b2cda6e 100644 --- a/backend/tests/test_gate.py +++ b/backend/tests/test_gate.py @@ -84,7 +84,7 @@ def __init__(self, connection_id: str) -> None: class _FakeConnections: @staticmethod - def list_all(): + async def list_all(): return [_FakeConnection("login-1"), _FakeConnection("login-2")] diff --git a/backend/tests/test_gate_receipt.py b/backend/tests/test_gate_receipt.py index 998b821e..aa1318aa 100644 --- a/backend/tests/test_gate_receipt.py +++ b/backend/tests/test_gate_receipt.py @@ -32,15 +32,15 @@ async def _call_through(options, func, *args, **kwargs): monkeypatch.setattr(DBOS, "run_step_async", _call_through) -def _reload(druks_db, run_id: str) -> Run: - druks_db.expire_all() - return druks_db.get(Run, run_id) +async def _reload(druks_db, run_id: str) -> Run: + druks_db.expunge_all() + return await druks_db.get(Run, run_id) async def test_answer_stamps_the_receipt_beside_the_gate_clear( druks_db, _direct_steps, monkeypatch ): - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-answer") + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-answer") async def _answer(topic, timeout_seconds): return {"action": "approve"} @@ -49,7 +49,7 @@ async def _answer(topic, timeout_seconds): payload = await _park(_ParkedWorkflow(run.id), "review", _ASK, ttl_seconds=1.0) assert payload == {"action": "approve"} - run = _reload(druks_db, run.id) + run = await _reload(druks_db, run.id) # The receipt is the round the answer cleared: the same stamp the park # wrote, which _GATE_CLEARED preserves on the row. assert run.input_requested_at @@ -59,7 +59,7 @@ async def _answer(topic, timeout_seconds): async def test_timeout_never_writes_the_receipt(druks_db, _direct_steps, monkeypatch): - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-timeout") + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-timeout") async def _lapse(topic, timeout_seconds): return None @@ -68,13 +68,13 @@ async def _lapse(topic, timeout_seconds): with pytest.raises(GateTimeout): await _park(_ParkedWorkflow(run.id), "review", _ASK, ttl_seconds=1.0) - run = _reload(druks_db, run.id) + run = await _reload(druks_db, run.id) assert not run.answer_parked_at assert run.input_requested_at async def test_cancel_never_writes_the_receipt(druks_db, _direct_steps, monkeypatch): - run = seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-cancel") + run = await seed_run(druks_db, kind=Summarize.kind, run_id="run-receipt-cancel") async def _cancelled(topic, timeout_seconds): raise DBOSWorkflowCancelledError(run.id) @@ -83,5 +83,5 @@ async def _cancelled(topic, timeout_seconds): with pytest.raises(DBOSWorkflowCancelledError): await _park(_ParkedWorkflow(run.id), "review", _ASK, ttl_seconds=1.0) - run = _reload(druks_db, run.id) + run = await _reload(druks_db, run.id) assert not run.answer_parked_at diff --git a/backend/tests/test_generic_subjects.py b/backend/tests/test_generic_subjects.py index 1d7cd2d3..e091cdd8 100644 --- a/backend/tests/test_generic_subjects.py +++ b/backend/tests/test_generic_subjects.py @@ -8,7 +8,7 @@ from druks.durable.datastructures import Subject from druks.durable.schemas import SubjectSummary from druks.models import StoredSubject -from druks.testing import seed_dbos_status +from druks.testing import asgi_client, seed_dbos_status from fastapi import APIRouter from fastapi.testclient import TestClient from pydantic import ValidationError @@ -30,8 +30,8 @@ def get_summary(self) -> _ThingSummary: return _ThingSummary(id=self.id, label=self.label, title=TITLES[self.id]) @classmethod - def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: - return [thing.get_summary() for thing in db_session().scalars(select(cls))] + async def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: + return [thing.get_summary() for thing in await db_session().scalars(select(cls))] class Ticket(Subject): @@ -39,7 +39,7 @@ class Ticket(Subject): separators a URL path is cut on.""" @classmethod - def get_for_subject_id(cls, subject_id: str) -> "Ticket | None": + async def get_for_subject_id(cls, subject_id: str) -> "Ticket | None": if "#" in subject_id: return cls(id=subject_id) return @@ -48,8 +48,8 @@ def get_summary(self) -> _ThingSummary: return _ThingSummary(id=self.id, label=self.label, title=self.id.rpartition("#")[2]) @classmethod - def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: - return [ticket.get_summary() for ticket in cls.list_open()] + async def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: + return [ticket.get_summary() for ticket in await cls.list_open()] CALLERS: list[str | None] = [] @@ -59,7 +59,7 @@ class Inbox(Subject): """A board scoped by who is asking.""" @classmethod - def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: + async def list_summaries(cls, account_id: str | None) -> list[_ThingSummary]: CALLERS.append(account_id) return [] @@ -68,7 +68,7 @@ class _ThingApp(App): name = "faketest" -def _seed_run( +async def _seed_run( session, *, subject_id, @@ -89,20 +89,20 @@ def _seed_run( failure=failure, ) session.add(run) - session.flush() - seed_dbos_status(session, run.id, state, subject={"type": subject_type, "id": subject_id}) + await session.flush() + await seed_dbos_status(session, run.id, state, subject={"type": subject_type, "id": subject_id}) return run -def _seed_call(session, run, *, agent, status="succeeded"): +async def _seed_call(session, run, *, agent, status="succeeded"): call = AgentCall(run_id=run.id, agent=agent, model="m", status=status, sandbox_host_id="h") session.add(call) - session.flush() + await session.flush() return call @pytest.fixture -def client(tmp_path: Path, druks_db, monkeypatch): +async def client(tmp_path: Path, druks_db, monkeypatch): # The real app mounts every app's routers before its catch-all 404, so the # fake app's router has to slot in there too — appending lands after the # catch-all and gets shadowed. Pulled back out on teardown; the app is a singleton. @@ -110,8 +110,8 @@ def client(tmp_path: Path, druks_db, monkeypatch): monkeypatch.setenv("DRUKS_DATA_DIR", str(tmp_path)) for subject_id in TITLES: - druks_db.merge(Thing(id=subject_id)) - druks_db.flush() + await druks_db.merge(Thing(id=subject_id)) + await druks_db.flush() app = configure_app_for_test(settings=make_settings(tmp_path)) holder = APIRouter() @@ -123,7 +123,7 @@ def client(tmp_path: Path, druks_db, monkeypatch): for route in reversed(holder.routes): app.router.routes.insert(catchall, route) try: - with TestClient(app) as test_client: + async with asgi_client(app) as test_client: yield test_client finally: for route in holder.routes: @@ -154,16 +154,16 @@ def test_a_summary_carries_the_subjects_own_label(druks_db): _ThingSummary(id=1, label=" ", title="First") -def test_status_aggregates_across_runs_and_timeline_spans_them(client: TestClient, druks_db): +async def test_status_aggregates_across_runs_and_timeline_spans_them(client: TestClient, druks_db): # Subject "1" lived across two runs: an earlier finished one and a current # running one. Status is the newest run's, and the timeline is every run, # oldest first, each carrying its own agent calls. - done = _seed_run(druks_db, subject_id="1", kind="faketest.prepare", state="finished") - _seed_call(druks_db, done, agent="prepare") - live = _seed_run(druks_db, subject_id="1", state="running") - _seed_call(druks_db, live, agent="implement", status="running") + done = await _seed_run(druks_db, subject_id="1", kind="faketest.prepare", state="finished") + await _seed_call(druks_db, done, agent="prepare") + live = await _seed_run(druks_db, subject_id="1", state="running") + await _seed_call(druks_db, live, agent="implement", status="running") - detail = client.get("/api/faketest/thing/1").json() + detail = (await client.get("/api/faketest/thing/1")).json() assert detail["summary"] == {"id": "1", "label": "thing 1", "title": "First"} assert detail["status"]["state"] == "running" assert [entry["kind"] for entry in detail["timeline"]] == ["faketest.prepare", "faketest.flow"] @@ -172,49 +172,49 @@ def test_status_aggregates_across_runs_and_timeline_spans_them(client: TestClien assert [c["agent"] for c in detail["timeline"][1]["agentCalls"]] == ["implement"] -def test_parked_run_surfaces_needs_you(client: TestClient, druks_db): - run = _seed_run( +async def test_parked_run_surfaces_needs_you(client: TestClient, druks_db): + run = await _seed_run( druks_db, subject_id="1", state="parked", input_gate="approve_plan", input_request={"label": "Approve the plan"}, ) - _seed_call(druks_db, run, agent="generate_plan") + await _seed_call(druks_db, run, agent="generate_plan") - detail = client.get("/api/faketest/thing/1").json() + detail = (await client.get("/api/faketest/thing/1")).json() assert detail["status"]["state"] == "parked" assert detail["status"]["gate"] == "approve_plan" parked = detail["timeline"][-1] assert parked["inputRequest"] == {"label": "Approve the plan"} -def test_status_carries_the_latest_run_failure(client: TestClient, druks_db): +async def test_status_carries_the_latest_run_failure(client: TestClient, druks_db): # A failed subject exposes its stop reason on the status, so a board can render # "why" without walking the timeline. An active or finished subject carries none. - _seed_run(druks_db, subject_id="1", state="failed", failure="profiler boom") + await _seed_run(druks_db, subject_id="1", state="failed", failure="profiler boom") - status = client.get("/api/faketest/thing/1").json()["status"] + status = (await client.get("/api/faketest/thing/1")).json()["status"] assert status["state"] == "failed" assert status["failure"] == "profiler boom" - _seed_run(druks_db, subject_id="2", state="running") - running = client.get("/api/faketest/thing/2").json()["status"] + await _seed_run(druks_db, subject_id="2", state="running") + running = (await client.get("/api/faketest/thing/2")).json()["status"] assert running["failure"] is None -def test_parked_board_row_skips_the_agent_call_query(client: TestClient, druks_db): +async def test_parked_board_row_skips_the_agent_call_query(client: TestClient, druks_db): # A parked row's status carries its gate ask, never its latest agent call, so # the per-subject status read must not load agent_calls — the board runs it # for every subject. - run = _seed_run( + run = await _seed_run( druks_db, subject_id="1", state="parked", input_gate="approve_plan", input_request={"label": "Approve the plan"}, ) - _seed_call(druks_db, run, agent="generate_plan") + await _seed_call(druks_db, run, agent="generate_plan") call_reads: list[str] = [] @@ -222,10 +222,10 @@ def record(conn, cursor, statement, parameters, context, executemany): if "agent_calls" in statement and statement.lstrip().upper().startswith("SELECT"): call_reads.append(statement) - engine = druks_db.get_bind() + engine = druks_db.bind.sync_connection event.listen(engine, "before_cursor_execute", record) try: - body = client.get("/api/faketest/thing").json() + body = (await client.get("/api/faketest/thing")).json() finally: event.remove(engine, "before_cursor_execute", record) @@ -234,11 +234,11 @@ def record(conn, cursor, statement, parameters, context, executemany): assert call_reads == [] -def test_list_returns_every_subject_with_status(client: TestClient, druks_db): - live = _seed_run(druks_db, subject_id="1", state="running") - _seed_call(druks_db, live, agent="implement", status="running") +async def test_list_returns_every_subject_with_status(client: TestClient, druks_db): + live = await _seed_run(druks_db, subject_id="1", state="running") + await _seed_call(druks_db, live, agent="implement", status="running") - body = client.get("/api/faketest/thing").json() + body = (await client.get("/api/faketest/thing")).json() rows = {row["summary"]["id"]: row for row in body["rows"]} assert rows["1"]["summary"]["title"] == "First" assert rows["1"]["status"]["state"] == "running" @@ -257,7 +257,7 @@ async def test_the_board_and_its_stream_hand_the_caller_to_list_summaries(druks_ token = current_account_id.set("acct-7") try: await endpoints["/inbox"]() - response = await endpoints["/inbox/stream"](engine=druks_db.get_bind()) + response = await endpoints["/inbox/stream"](engine=druks_db.bind) finally: current_account_id.reset(token) assert CALLERS == ["acct-7"] @@ -267,32 +267,36 @@ async def test_the_board_and_its_stream_hand_the_caller_to_list_summaries(druks_ assert CALLERS == ["acct-7", "acct-7"] -def test_unknown_subject_is_404(client: TestClient, druks_db): - assert client.get("/api/faketest/thing/nope").status_code == 404 +async def test_unknown_subject_is_404(client: TestClient, druks_db): + assert (await client.get("/api/faketest/thing/nope")).status_code == 404 # An id the subject could never wear misses the same way, row or no row. - assert client.get("/api/faketest/ticket/nope").status_code == 404 + assert (await client.get("/api/faketest/ticket/nope")).status_code == 404 -def test_an_id_spanning_separators_reaches_the_board_and_its_page(client: TestClient, druks_db): +async def test_an_id_spanning_separators_reaches_the_board_and_its_page( + client: TestClient, druks_db +): # A row-less subject's id is free text — "owner/repo#7" carries the path # separator and the fragment marker, and both reads still key on the whole id. - _seed_run(druks_db, subject_type="ticket", subject_id="owner/repo#7", state="parked") + await _seed_run(druks_db, subject_type="ticket", subject_id="owner/repo#7", state="parked") - board = client.get("/api/faketest/ticket").json() + board = (await client.get("/api/faketest/ticket")).json() assert [row["summary"]["id"] for row in board["rows"]] == ["owner/repo#7"] - detail = client.get("/api/faketest/ticket/owner/repo%237").json() + detail = (await client.get("/api/faketest/ticket/owner/repo%237")).json() assert detail["summary"] == {"id": "owner/repo#7", "label": "owner/repo#7", "title": "7"} assert detail["status"]["state"] == "parked" assert [entry["kind"] for entry in detail["timeline"]] == ["faketest.flow"] @pytest.mark.parametrize("path", ["thing/nope", "ticket/owner/nope"]) -def test_a_subjects_stream_wins_over_the_greedy_id_matcher(client: TestClient, druks_db, path): +async def test_a_subjects_stream_wins_over_the_greedy_id_matcher( + client: TestClient, druks_db, path +): # The id matcher spans separators, so ``/stream`` has to stay a suffix and not # get swallowed into the id — whatever shape the id is. A stream for a subject # that names nothing closes at once, which is what proves it got there. - response = client.get(f"/api/faketest/{path}/stream") + response = await client.get(f"/api/faketest/{path}/stream") assert response.status_code == 200 assert response.text == "" diff --git a/backend/tests/test_harness_auth.py b/backend/tests/test_harness_auth.py index 046159a3..9e6a4535 100644 --- a/backend/tests/test_harness_auth.py +++ b/backend/tests/test_harness_auth.py @@ -36,8 +36,10 @@ def _claude_payload(*, access="A0", refresh="R0", expires_at=None, extra=None) - return {"claudeAiOauth": block} -def _seed_claude(*, provider_email="op@example.com", **kwargs) -> HarnessConnection: - return connect_harness(ClaudeHarness, _claude_payload(**kwargs), provider_email=provider_email) +async def _seed_claude(*, provider_email="op@example.com", **kwargs) -> HarnessConnection: + return await connect_harness( + ClaudeHarness, _claude_payload(**kwargs), provider_email=provider_email + ) def _codex_payload(*, access=None, refresh="R0", account_id="acc-1", id_token="id-0") -> dict: @@ -48,8 +50,10 @@ def _codex_payload(*, access=None, refresh="R0", account_id="acc-1", id_token="i return {"auth_mode": "chatgpt", "OPENAI_API_KEY": None, "tokens": tokens} -def _seed_codex(*, provider_email="op@example.com", **kwargs) -> HarnessConnection: - return connect_harness(CodexHarness, _codex_payload(**kwargs), provider_email=provider_email) +async def _seed_codex(*, provider_email="op@example.com", **kwargs) -> HarnessConnection: + return await connect_harness( + CodexHarness, _codex_payload(**kwargs), provider_email=provider_email + ) def _resp(status: int, body: object) -> httpx.Response: @@ -83,44 +87,46 @@ async def fake_get(self, url, *, headers=None, **_kwargs): return calls -def test_claude_load_token(druks_db): - connection = _seed_claude(access="live", expires_at=_NOW + timedelta(hours=2)) +async def test_claude_load_token(druks_db): + connection = await _seed_claude(access="live", expires_at=_NOW + timedelta(hours=2)) token = ClaudeHarness.load_token(connection, now=_NOW) assert token.access_token == "live" assert token.subscription_type == "max" assert "user:profile" in token.scopes -def test_claude_load_token_expired(druks_db): - connection = _seed_claude(expires_at=_NOW - timedelta(hours=1)) +async def test_claude_load_token_expired(druks_db): + connection = await _seed_claude(expires_at=_NOW - timedelta(hours=1)) with pytest.raises(OAuthTokenError) as e: ClaudeHarness.load_token(connection, now=_NOW) assert e.value.tag == "token_expired" -def test_claude_load_token_no_access(druks_db): - connection = connect_harness(ClaudeHarness, {"claudeAiOauth": {"subscriptionType": "max"}}) +async def test_claude_load_token_no_access(druks_db): + connection = await connect_harness( + ClaudeHarness, {"claudeAiOauth": {"subscriptionType": "max"}} + ) with pytest.raises(OAuthTokenError) as e: ClaudeHarness.load_token(connection, now=_NOW) assert e.value.tag == "no_token" -def test_codex_load_token(druks_db): - connection = _seed_codex() +async def test_codex_load_token(druks_db): + connection = await _seed_codex() token = CodexHarness.load_token(connection, now=_NOW) assert "." in token.access_token assert token.account_id == "acc-1" -def test_codex_load_token_expired(druks_db): - connection = _seed_codex(access=_jwt(int((_NOW - timedelta(hours=1)).timestamp()))) +async def test_codex_load_token_expired(druks_db): + connection = await _seed_codex(access=_jwt(int((_NOW - timedelta(hours=1)).timestamp()))) with pytest.raises(OAuthTokenError) as e: CodexHarness.load_token(connection, now=_NOW) assert e.value.tag == "token_expired" async def test_claude_fresh_not_refreshed(monkeypatch, druks_db): - connection = _seed_claude(expires_at=_NOW + timedelta(hours=6)) + connection = await _seed_claude(expires_at=_NOW + timedelta(hours=6)) calls = _mock_post(monkeypatch, _resp(200, {})) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.action == "fresh" @@ -130,14 +136,14 @@ async def test_claude_fresh_not_refreshed(monkeypatch, druks_db): async def test_claude_stale_refreshes_and_persists(monkeypatch, druks_db): soon = _NOW + timedelta(minutes=30) - connection = _seed_claude(access="old", refresh="R0", expires_at=soon) + connection = await _seed_claude(access="old", refresh="R0", expires_at=soon) calls = _mock_post( monkeypatch, _resp(200, {"access_token": "new", "refresh_token": "R1", "expires_in": 28800}) ) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.action == "refreshed" assert calls[0]["json"]["refresh_token"] == "R0" - block = ClaudeHarness.get_credentials()["claudeAiOauth"] + block = (await ClaudeHarness.get_credentials())["claudeAiOauth"] assert block["accessToken"] == "new" assert block["refreshToken"] == "R1" assert block["scopes"] == ["user:profile"] # preserved @@ -146,44 +152,44 @@ async def test_claude_stale_refreshes_and_persists(monkeypatch, druks_db): async def test_claude_invalid_grant_drops_row(monkeypatch, druks_db): - connection = _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) _mock_post(monkeypatch, _resp(400, {"error": "invalid_grant"})) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.action == "failed" assert result.error == "invalid_grant" # A revoked lineage self-disconnects and commits inside the rotation — the # deletion never rides (or rolls back with) the tick's later commit. - assert not HarnessConnection.list_all() + assert not await HarnessConnection.list_all() with pytest.raises(HarnessNotConnectedError): - ClaudeHarness.get_credentials() + await ClaudeHarness.get_credentials() async def test_claude_network_error_keeps_row(monkeypatch, druks_db): - connection = _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) _mock_post(monkeypatch, httpx.ConnectError("boom")) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.error == "network" - assert ClaudeHarness.get_credentials()["claudeAiOauth"]["accessToken"] == "old" + assert (await ClaudeHarness.get_credentials())["claudeAiOauth"]["accessToken"] == "old" async def test_claude_http_500_keeps_row(monkeypatch, druks_db): - connection = _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) _mock_post(monkeypatch, _resp(500, "")) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.error == "http_500" - assert ClaudeHarness.get_credentials()["claudeAiOauth"]["accessToken"] == "old" + assert (await ClaudeHarness.get_credentials())["claudeAiOauth"]["accessToken"] == "old" async def test_claude_bad_response_keeps_row(monkeypatch, druks_db): - connection = _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) _mock_post(monkeypatch, _resp(200, "not json")) result = await ClaudeHarness.rotate_token(connection.id, now=_NOW) assert result.error == "bad_response" - assert ClaudeHarness.get_credentials()["claudeAiOauth"]["accessToken"] == "old" + assert (await ClaudeHarness.get_credentials())["claudeAiOauth"]["accessToken"] == "old" async def test_rotation_of_a_deleted_row_is_a_no_op(monkeypatch, druks_db): - connection = _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(access="old", expires_at=_NOW - timedelta(minutes=1)) connection_id = connection.id _mock_post(monkeypatch, _resp(400, {"error": "invalid_grant"})) await ClaudeHarness.rotate_token(connection_id, now=_NOW) @@ -197,7 +203,7 @@ async def test_rotation_of_a_deleted_row_is_a_no_op(monkeypatch, druks_db): async def test_claude_relogin_overwrite_picked_up(monkeypatch, druks_db): - connection = _seed_claude(refresh="R_NEW", expires_at=_NOW - timedelta(minutes=1)) + connection = await _seed_claude(refresh="R_NEW", expires_at=_NOW - timedelta(minutes=1)) calls = _mock_post( monkeypatch, _resp(200, {"access_token": "a", "refresh_token": "b", "expires_in": 100}) ) @@ -208,14 +214,14 @@ async def test_claude_relogin_overwrite_picked_up(monkeypatch, druks_db): async def test_codex_stale_refreshes_and_preserves(monkeypatch, druks_db): stale = _jwt(int((_NOW + timedelta(hours=1)).timestamp())) fresh = _jwt(int((_NOW + timedelta(days=10)).timestamp())) - connection = _seed_codex(access=stale, refresh="R0", account_id="acc-9") + connection = await _seed_codex(access=stale, refresh="R0", account_id="acc-9") calls = _mock_post( monkeypatch, _resp(200, {"access_token": fresh, "refresh_token": "R1", "id_token": "id-1"}) ) result = await CodexHarness.rotate_token(connection.id, now=_NOW) assert result.action == "refreshed" assert calls[0]["json"]["client_id"] == "app_EMoamEEZ73f0CkXaXp7hrann" - data = CodexHarness.get_credentials() + data = await CodexHarness.get_credentials() assert data["tokens"]["access_token"] == fresh assert data["tokens"]["refresh_token"] == "R1" assert data["tokens"]["id_token"] == "id-1" @@ -227,15 +233,15 @@ async def test_codex_stale_refreshes_and_preserves(monkeypatch, druks_db): async def test_codex_keeps_refresh_when_omitted(monkeypatch, druks_db): stale = _jwt(int((_NOW + timedelta(hours=1)).timestamp())) fresh = _jwt(int((_NOW + timedelta(days=10)).timestamp())) - connection = _seed_codex(access=stale, refresh="KEEP") + connection = await _seed_codex(access=stale, refresh="KEEP") _mock_post(monkeypatch, _resp(200, {"access_token": fresh})) await CodexHarness.rotate_token(connection.id, now=_NOW) - assert CodexHarness.get_credentials()["tokens"]["refresh_token"] == "KEEP" + assert (await CodexHarness.get_credentials())["tokens"]["refresh_token"] == "KEEP" async def test_codex_no_refresh_token(monkeypatch, druks_db): stale = _jwt(int((_NOW + timedelta(hours=1)).timestamp())) - connection = _seed_codex(refresh=None, access=stale) + connection = await _seed_codex(refresh=None, access=stale) calls = _mock_post(monkeypatch, _resp(200, {})) result = await CodexHarness.rotate_token(connection.id, now=_NOW) assert result.action == "no_refresh_token" @@ -243,13 +249,13 @@ async def test_codex_no_refresh_token(monkeypatch, druks_db): async def test_rotation_touches_only_the_addressed_row(monkeypatch, druks_db): - stale = _seed_claude( + stale = await _seed_claude( access="old", refresh="R0", expires_at=_NOW + timedelta(minutes=30), provider_email="a@example.com", ) - other = _seed_claude( + other = await _seed_claude( access="keep", refresh="RK", expires_at=_NOW + timedelta(minutes=30), @@ -261,24 +267,32 @@ async def test_rotation_touches_only_the_addressed_row(monkeypatch, druks_db): ) result = await ClaudeHarness.rotate_token(stale_id, now=_NOW) assert result.action == "refreshed" - assert dict(HarnessConnection.get(stale_id).payload)["claudeAiOauth"]["accessToken"] == "new" - assert dict(HarnessConnection.get(other_id).payload)["claudeAiOauth"]["accessToken"] == "keep" + assert ( + dict((await HarnessConnection.get(stale_id)).payload)["claudeAiOauth"]["accessToken"] + == "new" + ) + assert ( + dict((await HarnessConnection.get(other_id)).payload)["claudeAiOauth"]["accessToken"] + == "keep" + ) async def test_invalid_grant_drops_only_the_addressed_row(monkeypatch, druks_db): - kept = _seed_claude(access="d", expires_at=_NOW - timedelta(minutes=1)) - other = _seed_claude( + kept = await _seed_claude(access="d", expires_at=_NOW - timedelta(minutes=1)) + other = await _seed_claude( access="o", expires_at=_NOW - timedelta(minutes=1), provider_email="b@example.com" ) kept_id, other_id = kept.id, other.id _mock_post(monkeypatch, _resp(400, {"error": "invalid_grant"})) await ClaudeHarness.rotate_token(other_id, now=_NOW) - assert not HarnessConnection.get(other_id) - assert HarnessConnection.get(kept_id) + assert not await HarnessConnection.get(other_id) + assert await HarnessConnection.get(kept_id) async def test_rotation_stands_down_while_the_lock_is_held(monkeypatch, druks_db): - connection = _seed_claude(access="old", refresh="R0", expires_at=_NOW + timedelta(minutes=30)) + connection = await _seed_claude( + access="old", refresh="R0", expires_at=_NOW + timedelta(minutes=30) + ) calls = _mock_post( monkeypatch, _resp(200, {"access_token": "new", "refresh_token": "R1", "expires_in": 100}) ) @@ -291,7 +305,9 @@ async def test_rotation_stands_down_while_the_lock_is_held(monkeypatch, druks_db async def test_rotation_lock_is_released_after_refresh(monkeypatch, druks_db): - connection = _seed_claude(access="old", refresh="R0", expires_at=_NOW + timedelta(minutes=30)) + connection = await _seed_claude( + access="old", refresh="R0", expires_at=_NOW + timedelta(minutes=30) + ) _mock_post( monkeypatch, _resp(200, {"access_token": "new", "refresh_token": "R1", "expires_in": 100}) ) @@ -299,54 +315,54 @@ async def test_rotation_lock_is_released_after_refresh(monkeypatch, druks_db): assert not await druks.redis.get_client().get(f"druks:harness:refresh:{connection.id}") -def test_disconnect_removes_only_the_addressed_login(druks_db): - mine = _seed_claude(provider_email="a@example.com") - other = _seed_claude(provider_email="b@example.com") +async def test_disconnect_removes_only_the_addressed_login(druks_db): + mine = await _seed_claude(provider_email="a@example.com") + other = await _seed_claude(provider_email="b@example.com") - mine.delete() + await mine.delete() - assert HarnessConnection.get(other.id) + assert await HarnessConnection.get(other.id) # The fallback account (the first) has no claude connection left; another # account's credential never leaks into execution. with pytest.raises(HarnessNotConnectedError): - ClaudeHarness.get_credentials() + await ClaudeHarness.get_credentials() -def test_reconnect_restores_execution(druks_db): - mine = _seed_claude(provider_email="a@example.com") - mine.delete() +async def test_reconnect_restores_execution(druks_db): + mine = await _seed_claude(provider_email="a@example.com") + await mine.delete() with pytest.raises(HarnessNotConnectedError): - ClaudeHarness.get_credentials() + await ClaudeHarness.get_credentials() - _seed_claude(access="fresh", provider_email="a@example.com") - assert ClaudeHarness.get_credentials()["claudeAiOauth"]["accessToken"] == "fresh" + await _seed_claude(access="fresh", provider_email="a@example.com") + assert (await ClaudeHarness.get_credentials())["claudeAiOauth"]["accessToken"] == "fresh" -def test_connect_scopes_rows_by_harness_and_account(druks_db): - claude_row = _seed_claude(provider_email="a@example.com") - codex_row = _seed_codex(provider_email="a@example.com") - other = _seed_claude(provider_email="b@example.com") +async def test_connect_scopes_rows_by_harness_and_account(druks_db): + claude_row = await _seed_claude(provider_email="a@example.com") + codex_row = await _seed_codex(provider_email="a@example.com") + other = await _seed_claude(provider_email="b@example.com") assert len({claude_row.id, codex_row.id, other.id}) == 3 assert claude_row.account_id == codex_row.account_id # same person, one account assert other.account_id != claude_row.account_id - assert Account.get_for_username("a@example.com").id == claude_row.account_id + assert (await Account.get_for_username("a@example.com")).id == claude_row.account_id # The first account adopted the execution fallback. - assert UserSettings.get().fallback_account_id == claude_row.account_id + assert (await UserSettings.get()).fallback_account_id == claude_row.account_id -def test_reconnect_updates_the_existing_login_in_place(druks_db): - row = _seed_claude(access="old", provider_email="a@example.com") +async def test_reconnect_updates_the_existing_login_in_place(druks_db): + row = await _seed_claude(access="old", provider_email="a@example.com") # Same email, different case — citext matches it to the existing account, # so the reconnect updates that one connection rather than making a second. - again = _seed_claude(access="new", provider_email="A@Example.com") + again = await _seed_claude(access="new", provider_email="A@Example.com") assert again.id == row.id assert dict(again.payload)["claudeAiOauth"]["accessToken"] == "new" assert again.provider_email == "A@Example.com" # stored as last given async def test_claude_fetch_usage_success(monkeypatch, druks_db): - connection = _seed_claude(access="tok", expires_at=_NOW + timedelta(hours=2)) + connection = await _seed_claude(access="tok", expires_at=_NOW + timedelta(hours=2)) body = { "five_hour": {"utilization": 16.0, "resets_at": "2026-06-04T23:19:59+00:00"}, "seven_day": {"utilization": 48.0, "resets_at": "2026-06-07T16:00:00+00:00"}, @@ -362,7 +378,7 @@ async def test_claude_fetch_usage_success(monkeypatch, druks_db): async def test_claude_fetch_usage_http_error(monkeypatch, druks_db): - connection = _seed_claude(access="tok", expires_at=_NOW + timedelta(hours=2)) + connection = await _seed_claude(access="tok", expires_at=_NOW + timedelta(hours=2)) _mock_get(monkeypatch, _resp(403, {"error": "x"})) parsed = await ClaudeHarness.fetch_usage(connection, now=_NOW) assert parsed.ok is False @@ -371,7 +387,7 @@ async def test_claude_fetch_usage_http_error(monkeypatch, druks_db): async def test_fetch_usage_without_a_token_skips_http(monkeypatch, druks_db): # The connection exists but its payload carries no access token — never fetch. - connection = connect_harness(ClaudeHarness, {"claudeAiOauth": {}}) + connection = await connect_harness(ClaudeHarness, {"claudeAiOauth": {}}) calls = _mock_get(monkeypatch, _resp(200, {})) parsed = await ClaudeHarness.fetch_usage(connection, now=_NOW) assert parsed.ok is False @@ -380,7 +396,7 @@ async def test_fetch_usage_without_a_token_skips_http(monkeypatch, druks_db): async def test_codex_fetch_usage_success(monkeypatch, druks_db): - connection = _seed_codex(account_id="acc-7") + connection = await _seed_codex(account_id="acc-7") body = { "plan_type": "pro", "rate_limit": { @@ -405,20 +421,20 @@ async def test_codex_fetch_usage_success(monkeypatch, druks_db): assert calls[0]["headers"]["ChatGPT-Account-Id"] == "acc-7" -def test_render_credentials_file_serializes_stored_payload(druks_db): - connection = _seed_claude(access="tok", refresh="R0") - rendered = ClaudeHarness.render_credentials_file(connection.id) +async def test_render_credentials_file_serializes_stored_payload(druks_db): + connection = await _seed_claude(access="tok", refresh="R0") + rendered = await ClaudeHarness.render_credentials_file(connection.id) assert json.loads(rendered)["claudeAiOauth"]["accessToken"] == "tok" -def test_render_credentials_file_raises_when_not_connected(druks_db): +async def test_render_credentials_file_raises_when_not_connected(druks_db): # No selection and no fallback connection at all. with pytest.raises(HarnessNotConnectedError, match="connect it in Settings"): - ClaudeHarness.render_credentials_file() + await ClaudeHarness.render_credentials_file() -def test_claude_builder_puts_db_credentials_on_the_bundle(druks_db): - _seed_claude(access="live", refresh="R0") +async def test_claude_builder_puts_db_credentials_on_the_bundle(druks_db): + await _seed_claude(access="live", refresh="R0") sandbox = SandboxSettings( service_url="x", service_token="x", @@ -427,14 +443,14 @@ def test_claude_builder_puts_db_credentials_on_the_bundle(druks_db): claude_config_dir=Path("/home/agent/.claude"), codex_config_dir=Path("/home/agent/.codex"), ) - bundle = _claude_credentials(sandbox, github_token=None) + bundle = await _claude_credentials(sandbox, github_token=None) assert json.loads(bundle.claude_credentials)["claudeAiOauth"]["accessToken"] == "live" assert bundle.codex_credentials is None -def test_credentials_builders_carry_global_instructions(druks_db): - _seed_claude() - _seed_codex() +async def test_credentials_builders_carry_global_instructions(druks_db): + await _seed_claude() + await _seed_codex() claude_config_dir = Path("/home/agent/.claude") codex_config_dir = Path("/home/agent/.codex") sandbox = SandboxSettings( @@ -446,8 +462,8 @@ def test_credentials_builders_carry_global_instructions(druks_db): codex_config_dir=codex_config_dir, ) - claude_credentials = _claude_credentials(sandbox, github_token=None) - codex_credentials = CodexHarness( + claude_credentials = await _claude_credentials(sandbox, github_token=None) + codex_credentials = await CodexHarness( model=CodexHarness.default_model, fast_mode=False, effort=None, @@ -462,11 +478,11 @@ def test_credentials_builders_carry_global_instructions(druks_db): ) -def test_no_config_dir_ships_credential_only(druks_db): +async def test_no_config_dir_ships_credential_only(druks_db): # No local config dir for the CLI => nothing of the host's config/plugins # reaches the sandbox — but the DB credential still ships: connection state # alone decides whether a harness can run. - _seed_claude(access="live") + await _seed_claude(access="live") sandbox = SandboxSettings( service_url="x", service_token="x", @@ -475,14 +491,14 @@ def test_no_config_dir_ships_credential_only(druks_db): claude_config_dir=None, codex_config_dir=None, ) - bundle = _claude_credentials(sandbox, github_token="gh") + bundle = await _claude_credentials(sandbox, github_token="gh") assert json.loads(bundle.claude_credentials)["claudeAiOauth"]["accessToken"] == "live" assert bundle.extra_config_files == () assert bundle.extra_config_dirs == () assert bundle.github_token == "gh" -def test_claude_builder_raises_when_not_connected(druks_db): +async def test_claude_builder_raises_when_not_connected(druks_db): sandbox = SandboxSettings( service_url="x", service_token="x", @@ -492,49 +508,49 @@ def test_claude_builder_raises_when_not_connected(druks_db): codex_config_dir=None, ) with pytest.raises(HarnessNotConnectedError, match="claude is not connected"): - _claude_credentials(sandbox, github_token=None) + await _claude_credentials(sandbox, github_token=None) -def test_lookup_prefers_the_accounts_own_connection(druks_db): - fallback = _seed_claude(provider_email="a@example.com") # a@ adopts the fallback - own = _seed_claude(provider_email="b@example.com") +async def test_lookup_prefers_the_accounts_own_connection(druks_db): + fallback = await _seed_claude(provider_email="a@example.com") # a@ adopts the fallback + own = await _seed_claude(provider_email="b@example.com") - assert HarnessConnection.lookup("claude", own.account_id).id == own.id - assert HarnessConnection.lookup("claude", fallback.account_id).id == fallback.id + assert (await HarnessConnection.lookup("claude", own.account_id)).id == own.id + assert (await HarnessConnection.lookup("claude", fallback.account_id)).id == fallback.id -def test_lookup_falls_back(druks_db): - fallback = _seed_claude(provider_email="a@example.com") - codex_only = _seed_codex(provider_email="b@example.com") +async def test_lookup_falls_back(druks_db): + fallback = await _seed_claude(provider_email="a@example.com") + codex_only = await _seed_codex(provider_email="b@example.com") # An account with no claude connection, and no account at all. - assert HarnessConnection.lookup("claude", codex_only.account_id).id == fallback.id - assert HarnessConnection.lookup("claude", None).id == fallback.id + assert (await HarnessConnection.lookup("claude", codex_only.account_id)).id == fallback.id + assert (await HarnessConnection.lookup("claude", None)).id == fallback.id -def test_lookup_without_any_connection_raises(druks_db): - _seed_codex(provider_email="a@example.com") # the fallback account has codex only +async def test_lookup_without_any_connection_raises(druks_db): + await _seed_codex(provider_email="a@example.com") # the fallback account has codex only with pytest.raises(HarnessNotConnectedError, match="connect it in Settings"): - HarnessConnection.lookup("claude", None) + await HarnessConnection.lookup("claude", None) -def test_render_credentials_file_renders_only_the_selected_login(druks_db): - mine = _seed_claude(access="mine-token", provider_email="a@example.com") - other = _seed_claude(access="other-token", provider_email="b@example.com") +async def test_render_credentials_file_renders_only_the_selected_login(druks_db): + mine = await _seed_claude(access="mine-token", provider_email="a@example.com") + other = await _seed_claude(access="other-token", provider_email="b@example.com") - rendered = json.loads(ClaudeHarness.render_credentials_file(other.id)) + rendered = json.loads(await ClaudeHarness.render_credentials_file(other.id)) assert rendered["claudeAiOauth"]["accessToken"] == "other-token" assert "mine-token" not in json.dumps(rendered) - rendered = json.loads(ClaudeHarness.render_credentials_file(mine.id)) + rendered = json.loads(await ClaudeHarness.render_credentials_file(mine.id)) assert rendered["claudeAiOauth"]["accessToken"] == "mine-token" -def test_render_credentials_file_for_a_deleted_connection_raises(druks_db): - _seed_claude(provider_email="a@example.com") # the surviving fallback - gone = _seed_claude(provider_email="b@example.com") +async def test_render_credentials_file_for_a_deleted_connection_raises(druks_db): + await _seed_claude(provider_email="a@example.com") # the surviving fallback + gone = await _seed_claude(provider_email="b@example.com") gone_id = gone.id - gone.delete() + await gone.delete() # A disconnect between selection and render fails the call — it must never # fall through to another account's payload. with pytest.raises(HarnessNotConnectedError, match="removed"): - ClaudeHarness.render_credentials_file(gone_id) + await ClaudeHarness.render_credentials_file(gone_id) diff --git a/backend/tests/test_harness_login_persistence.py b/backend/tests/test_harness_login_persistence.py index 0bf45d86..32e34cc5 100644 --- a/backend/tests/test_harness_login_persistence.py +++ b/backend/tests/test_harness_login_persistence.py @@ -7,6 +7,7 @@ from druks.harnesses.models import HarnessConnection from druks.testing import init_db from sqlalchemy import create_engine, text +from sqlalchemy.ext.asyncio import create_async_engine # The credential store's whole job is to persist a rotated credential dict # through a real commit. The rollback-based suite can't verify that — its identity @@ -31,42 +32,44 @@ def _pg_up() -> bool: @pytest.fixture -def engine(): +async def engine(): admin = psycopg.connect(f"{PG_BASE}/postgres", autocommit=True) admin.execute(f"DROP DATABASE IF EXISTS {DB}") admin.execute(f"CREATE DATABASE {DB}") admin.close() - created = create_engine(URL) - init_db(created) + schema_engine = create_engine(URL) + init_db(schema_engine) + schema_engine.dispose() + created = create_async_engine(URL) configure_session(created) try: yield created finally: - created.dispose() + await created.dispose() -def _committed(engine, work): +async def _committed(engine, work): session = get_session(engine) db_session.registry.set(session) try: - result = work() - session.commit() + result = await work() + await session.commit() return result finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() -def _connect(payload: dict) -> str: +async def _connect(payload: dict) -> str: from druks.accounts.models import Account from druks.user_settings.models import UserSettings - account = Account.get_or_create("op@example.com") - settings = UserSettings.get() + account = await Account.get_or_create("op@example.com") + settings = await UserSettings.get() if not settings.fallback_account_id: - settings.set_fallback_account(account.id) - row = HarnessConnection.connect( + await settings.set_fallback_account(account.id) + row = await HarnessConnection.connect( harness="claude", account=account, payload=payload, @@ -76,37 +79,46 @@ def _connect(payload: dict) -> str: return row.id -def test_rotation_persists_new_payload_across_sessions(engine): +async def test_rotation_persists_new_payload_across_sessions(engine): # Connect, then rotate the payload the way rotate_token does (plain-dict # copy, edit, whole-value update), then read it back from a fresh session — # commit + new session proves the edit reached the DB, not just the # in-memory object. - connection_id = _committed( - engine, lambda: _connect({"claudeAiOauth": {"accessToken": "old", "refreshToken": "R0"}}) - ) + async def connect_old(): + return await _connect({"claudeAiOauth": {"accessToken": "old", "refreshToken": "R0"}}) + + connection_id = await _committed(engine, connect_old) - def rotate_in_place(): - row = HarnessConnection.get(connection_id) + async def rotate_in_place(): + row = await HarnessConnection.get(connection_id) data = dict(row.payload) data["claudeAiOauth"]["accessToken"] = "new" - row.update_payload(data, expires_at=None) + await row.update_payload(data, expires_at=None) - _committed(engine, rotate_in_place) + await _committed(engine, rotate_in_place) - block = _committed( - engine, lambda: dict(HarnessConnection.get(connection_id).payload)["claudeAiOauth"] - ) + async def read_back(): + row = await HarnessConnection.get(connection_id) + return dict(row.payload)["claudeAiOauth"] + + block = await _committed(engine, read_back) assert block["accessToken"] == "new" -def test_payload_is_ciphertext_at_rest(engine): - _committed(engine, lambda: _connect({"claudeAiOauth": {"accessToken": "supersecret"}})) +async def test_payload_is_ciphertext_at_rest(engine): + async def connect_secret(): + return await _connect({"claudeAiOauth": {"accessToken": "supersecret"}}) + + await _committed(engine, connect_secret) - with engine.connect() as connection: - stored = connection.execute(text("SELECT payload FROM harness_logins")).scalar_one() + async with engine.connect() as connection: + stored = (await connection.execute(text("SELECT payload FROM harness_logins"))).scalar_one() raw = bytes(stored) assert b"supersecret" not in raw assert b"claudeAiOauth" not in raw - block = _committed(engine, lambda: ClaudeHarness.get_credentials()["claudeAiOauth"]) + async def read_credentials(): + return (await ClaudeHarness.get_credentials())["claudeAiOauth"] + + block = await _committed(engine, read_credentials) assert block["accessToken"] == "supersecret" diff --git a/backend/tests/test_harness_model_routing.py b/backend/tests/test_harness_model_routing.py index 192edc38..2e3fecba 100644 --- a/backend/tests/test_harness_model_routing.py +++ b/backend/tests/test_harness_model_routing.py @@ -10,35 +10,35 @@ from fastapi import HTTPException -def test_shipped_tuple_fallback_routes_shipped_models(druks_db): - assert get_harness_for_model("claude-opus-4-7") is ClaudeHarness - assert get_harness_for_model("gpt-5.5") is CodexHarness +async def test_shipped_tuple_fallback_routes_shipped_models(druks_db): + assert await get_harness_for_model("claude-opus-4-7") is ClaudeHarness + assert await get_harness_for_model("gpt-5.5") is CodexHarness -def test_fetched_list_routes_provider_models(druks_db): - HarnessSettings.require("claude").models_fetched = [ +async def test_fetched_list_routes_provider_models(druks_db): + (await HarnessSettings.require("claude")).models_fetched = [ {"id": "claude-fable-5", "label": "Claude Fable 5"} ] - druks_db.flush() + await druks_db.flush() - assert get_harness_for_model("claude-fable-5") is ClaudeHarness + assert await get_harness_for_model("claude-fable-5") is ClaudeHarness -def test_bare_harness_name_routes(druks_db): - assert get_harness_for_model("claude") is ClaudeHarness - assert get_harness_for_model("codex") is CodexHarness +async def test_bare_harness_name_routes(druks_db): + assert await get_harness_for_model("claude") is ClaudeHarness + assert await get_harness_for_model("codex") is CodexHarness -def test_unknown_model_raises_harness_error(druks_db): +async def test_unknown_model_raises_harness_error(druks_db): with pytest.raises(HarnessError): - get_harness_for_model("llama-3-70b") + await get_harness_for_model("llama-3-70b") with pytest.raises(HarnessError): - get_harness_for_model("claude-opus-99") + await get_harness_for_model("claude-opus-99") -def test_settings_reject_model_missing_from_lists_returns_422(druks_db): +async def test_settings_reject_model_missing_from_lists_returns_422(druks_db): with pytest.raises(HTTPException) as error: - _validate_model("llama-3-70b") + await _validate_model("llama-3-70b") assert error.value.status_code == 422 assert error.value.detail == "No installed harness runs model 'llama-3-70b'." @@ -52,8 +52,8 @@ def test_settings_reject_model_missing_from_lists_returns_422(druks_db): ], ) async def test_settings_reject_invalid_model_returns_422(druks_db, model, detail): - account = Account.get_or_create("op@example.com") - settings = HarnessSettings.require("claude") + account = await Account.get_or_create("op@example.com") + settings = await HarnessSettings.require("claude") original_model = settings.model with pytest.raises(HTTPException) as error: diff --git a/backend/tests/test_harness_reasoning_flags.py b/backend/tests/test_harness_reasoning_flags.py index 83fde3dc..693db39f 100644 --- a/backend/tests/test_harness_reasoning_flags.py +++ b/backend/tests/test_harness_reasoning_flags.py @@ -7,11 +7,11 @@ @pytest.fixture(autouse=True) -def _connected_harnesses(druks_db): +async def _connected_harnesses(druks_db): # build_invocation renders each credential bundle from the DB row and # raises when that harness isn't connected. - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) - connect_harness(CodexHarness, {"tokens": {"access_token": "t"}}) + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "t"}}) + await connect_harness(CodexHarness, {"tokens": {"access_token": "t"}}) def _sandbox_config(): @@ -29,7 +29,7 @@ def _sandbox_config(): ) -def test_claude_build_invocation_carries_every_flag(): +async def test_claude_build_invocation_carries_every_flag(): """Flag-drop guard: moving argv construction is exactly how CLI flags got silently lost before — assert the full surface.""" import shlex @@ -38,7 +38,7 @@ def test_claude_build_invocation_carries_every_flag(): schema = {"type": "object"} server = McpServer(name="github", url="https://api.example/mcp/", bearer_token_env_var="TOK") - inv = ClaudeHarness( + inv = await ClaudeHarness( model="claude-x", fast_mode=True, effort="high", @@ -90,12 +90,12 @@ def test_claude_build_invocation_carries_every_flag(): assert inv.extra_artifact_filenames == ("debug.log", "session.jsonl") -def test_codex_build_invocation_carries_every_flag(): +async def test_codex_build_invocation_carries_every_flag(): """Flag-drop guard, codex side.""" from druks.sandbox.datastructures import McpServer server = McpServer(name="github", url="https://api.example/mcp/", bearer_token_env_var="TOK") - inv = CodexHarness( + inv = await CodexHarness( model=_CODEX_MODEL, fast_mode=True, effort="high", diff --git a/backend/tests/test_identity.py b/backend/tests/test_identity.py index bef7502d..5e5503fd 100644 --- a/backend/tests/test_identity.py +++ b/backend/tests/test_identity.py @@ -73,8 +73,8 @@ def _connect( ) -def _all_accounts() -> list[Account]: - return list(database.db_session().scalars(select(Account))) +async def _all_accounts() -> list[Account]: + return list(await database.db_session().scalars(select(Account))) def _mock_exchange_codex(monkeypatch, *, email: str): @@ -96,7 +96,7 @@ async def fake_post(self, url, *, json=None, data=None, **_kwargs): # --- header mode ----------------------------------------------------------- -def test_header_mode_requires_exactly_one_nonblank_assertion(tmp_path, druks_db): +async def test_header_mode_requires_exactly_one_nonblank_assertion(tmp_path, druks_db): with _header_client(tmp_path) as client: assert client.get("/api/auth/me").status_code == 401 assert client.get("/api/settings").status_code == 401 @@ -106,10 +106,10 @@ def test_header_mode_requires_exactly_one_nonblank_assertion(tmp_path, druks_db) ) assert two.status_code == 401 # Rejection never enrolls anyone. - assert {account.username for account in _all_accounts()} == {"system"} + assert {account.username for account in await _all_accounts()} == {"system"} -def test_an_asserted_email_open_enrolls_once_across_case_variants(tmp_path, druks_db): +async def test_an_asserted_email_open_enrolls_once_across_case_variants(tmp_path, druks_db): with _header_client(tmp_path) as client: first = client.get("/api/auth/me", headers={HEADER: " Op@Example.com "}) assert first.status_code == 200 @@ -120,21 +120,25 @@ def test_an_asserted_email_open_enrolls_once_across_case_variants(tmp_path, druk again = client.get("/api/auth/me", headers={HEADER: "op@example.COM"}) assert again.json()["account"]["id"] == body["account"]["id"] - assert len(Account.list_non_system()) == 1 + assert len(await Account.list_non_system()) == 1 -def test_get_or_create_losing_the_insert_race_still_converges(druks_db, monkeypatch): - existing = Account.get_or_create("race@example.com") +async def test_get_or_create_losing_the_insert_race_still_converges(druks_db, monkeypatch): + existing = await Account.get_or_create("race@example.com") # Simulate losing the read-then-insert race: the pre-read misses, the # INSERT hits ON CONFLICT DO NOTHING, the canonical lookup converges. - monkeypatch.setattr(Account, "get_for_username", classmethod(lambda cls, username: None)) - assert Account.get_or_create("Race@example.com").id == existing.id - assert len(Account.list_non_system()) == 1 + async def _miss(cls, username): + return None -def test_a_valid_pat_wins_over_a_conflicting_header(tmp_path, druks_db): - agent = Account.get_or_create("agent@example.com") - _, token = PersonalAccessToken.create(account_id=agent.id, name="agent") + monkeypatch.setattr(Account, "get_for_username", classmethod(_miss)) + assert (await Account.get_or_create("Race@example.com")).id == existing.id + assert len(await Account.list_non_system()) == 1 + + +async def test_a_valid_pat_wins_over_a_conflicting_header(tmp_path, druks_db): + agent = await Account.get_or_create("agent@example.com") + _, token = await PersonalAccessToken.create(account_id=agent.id, name="agent") with _header_client(tmp_path) as client: response = client.get( "/api/auth/me", @@ -143,11 +147,11 @@ def test_a_valid_pat_wins_over_a_conflicting_header(tmp_path, druks_db): assert response.status_code == 200 assert response.json()["account"]["username"] == "agent@example.com" # The losing assertion never enrolled. - assert not Account.get_for_username("op@example.com") + assert not await Account.get_for_username("op@example.com") -def test_onboarding_clears_once_the_account_has_a_connection(tmp_path, druks_db): - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) +async def test_onboarding_clears_once_the_account_has_a_connection(tmp_path, druks_db): + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) with _header_client(tmp_path) as client: body = client.get("/api/auth/me", headers={HEADER: "op@example.com"}).json() assert body["onboardingRequired"] is False @@ -156,14 +160,14 @@ def test_onboarding_clears_once_the_account_has_a_connection(tmp_path, druks_db) # --- none mode ------------------------------------------------------------- -def test_none_mode_ignores_a_present_identity_header(tmp_path, druks_db): - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) +async def test_none_mode_ignores_a_present_identity_header(tmp_path, druks_db): + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) with _client(tmp_path) as client: body = client.get("/api/auth/me", headers={HEADER: "intruder@example.com"}).json() assert body["authMode"] == "none" assert body["account"]["username"] == "op@example.com" # Never open-enrolls in none mode. - assert not Account.get_for_username("intruder@example.com") + assert not await Account.get_for_username("intruder@example.com") def test_none_zero_reads_as_setup(tmp_path, druks_db): @@ -174,8 +178,8 @@ def test_none_zero_reads_as_setup(tmp_path, druks_db): assert client.get("/api/settings").status_code == 409 -def test_none_one_resolves_the_operator(tmp_path, druks_db): - connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) +async def test_none_one_resolves_the_operator(tmp_path, druks_db): + await connect_harness(ClaudeHarness, {"claudeAiOauth": {"accessToken": "x"}}) with _client(tmp_path) as client: body = client.get("/api/auth/me").json() assert body["account"]["username"] == "op@example.com" @@ -183,20 +187,20 @@ def test_none_one_resolves_the_operator(tmp_path, druks_db): assert client.get("/api/settings").status_code == 200 -def test_none_multi_refuses_requests_and_startup(tmp_path, druks_db): - Account.get_or_create("one@example.com") - Account.get_or_create("two@example.com") +async def test_none_multi_refuses_requests_and_startup(tmp_path, druks_db): + await Account.get_or_create("one@example.com") + await Account.get_or_create("two@example.com") with _client(tmp_path) as client: assert client.get("/api/settings").status_code == 503 # The startup validator runs the same check and refuses boot. with pytest.raises(AuthConfigurationError): - resolve_single_operator() + await resolve_single_operator() # --- connection flow ------------------------------------------------------- -def test_none_zero_setup_flow_creates_the_operator(tmp_path, monkeypatch, druks_db): +async def test_none_zero_setup_flow_creates_the_operator(tmp_path, monkeypatch, druks_db): with _client(tmp_path) as client: response = _connect(client, monkeypatch, email="me@example.com") assert response.status_code == 200 @@ -206,12 +210,14 @@ def test_none_zero_setup_flow_creates_the_operator(tmp_path, monkeypatch, druks_ body = client.get("/api/auth/me").json() assert body["account"]["username"] == "me@example.com" assert body["onboardingRequired"] is False - account = Account.get_for_username("me@example.com") - assert UserSettings.get().fallback_account_id == account.id - assert HarnessConnection.get_for_account("claude", account.id) + account = await Account.get_for_username("me@example.com") + assert (await UserSettings.get()).fallback_account_id == account.id + assert await HarnessConnection.get_for_account("claude", account.id) -def test_concurrent_setup_completions_with_one_email_converge(tmp_path, monkeypatch, druks_db): +async def test_concurrent_setup_completions_with_one_email_converge( + tmp_path, monkeypatch, druks_db +): with _client(tmp_path) as client: # Both flows start while zero accounts exist — both unbound. first = client.post("/api/harnesses/claude/connection/start") @@ -232,11 +238,11 @@ def test_concurrent_setup_completions_with_one_email_converge(tmp_path, monkeypa ).status_code == 200 ) - assert len(Account.list_non_system()) == 1 - assert len(HarnessConnection.list_all()) == 2 + assert len(await Account.list_non_system()) == 1 + assert len(await HarnessConnection.list_all()) == 2 -def test_a_stale_unbound_completion_attaches_to_the_operator(tmp_path, monkeypatch, druks_db): +async def test_a_stale_unbound_completion_attaches_to_the_operator(tmp_path, monkeypatch, druks_db): with _client(tmp_path) as client: # Both flows start while zero accounts exist; the first completion # creates the operator, so the second — a different provider email — @@ -257,14 +263,14 @@ def test_a_stale_unbound_completion_attaches_to_the_operator(tmp_path, monkeypat assert completed.status_code == 200 assert completed.json()["username"] == "a@example.com" assert client.get("/api/settings").status_code == 200 - operator = Account.get_for_username("a@example.com") - assert len(Account.list_non_system()) == 1 - codex_connection = HarnessConnection.get_for_account("codex", operator.id) + operator = await Account.get_for_username("a@example.com") + assert len(await Account.list_non_system()) == 1 + codex_connection = await HarnessConnection.get_for_account("codex", operator.id) # The capability keeps its own provider identity; it never rekeys the account. assert codex_connection.provider_email == "b@example.com" -def test_a_connect_survives_a_failed_model_refresh(tmp_path, monkeypatch, druks_db): +async def test_a_connect_survives_a_failed_model_refresh(tmp_path, monkeypatch, druks_db): async def _refresh_boom(self, connection): raise RuntimeError("picker flush failed") @@ -274,11 +280,13 @@ async def _refresh_boom(self, connection): with _client(tmp_path) as client: response = _connect(client, monkeypatch, email="me@example.com") assert response.status_code == 200 - account = Account.get_for_username("me@example.com") - assert HarnessConnection.get_for_account("claude", account.id) + account = await Account.get_for_username("me@example.com") + assert await HarnessConnection.get_for_account("claude", account.id) -def test_a_bound_connect_cannot_complete_under_another_operator(tmp_path, monkeypatch, druks_db): +async def test_a_bound_connect_cannot_complete_under_another_operator( + tmp_path, monkeypatch, druks_db +): with _header_client(tmp_path) as client: start = client.post( "/api/harnesses/claude/connection/start", headers={HEADER: "alice@example.com"} @@ -291,14 +299,14 @@ def test_a_bound_connect_cannot_complete_under_another_operator(tmp_path, monkey ) assert response.status_code == 422 assert "different operator" in response.json()["detail"] - assert not any(row.harness == "claude" for row in HarnessConnection.list_all()) + assert not any(row.harness == "claude" for row in await HarnessConnection.list_all()) -def test_first_connection_claims_the_fallback_slot_once(tmp_path, monkeypatch, druks_db): +async def test_first_connection_claims_the_fallback_slot_once(tmp_path, monkeypatch, druks_db): with _header_client(tmp_path) as client: _connect(client, monkeypatch, email="seat@corp.com", headers={HEADER: "first@example.com"}) - first = Account.get_for_username("first@example.com") - assert UserSettings.get().fallback_account_id == first.id + first = await Account.get_for_username("first@example.com") + assert (await UserSettings.get()).fallback_account_id == first.id _connect( client, monkeypatch, @@ -307,10 +315,12 @@ def test_first_connection_claims_the_fallback_slot_once(tmp_path, monkeypatch, d headers={HEADER: "second@example.com"}, ) # The fallback stays with the first operator. - assert UserSettings.get().fallback_account_id == first.id + assert (await UserSettings.get()).fallback_account_id == first.id -def test_reconnect_records_provider_email_but_keeps_the_operator(tmp_path, monkeypatch, druks_db): +async def test_reconnect_records_provider_email_but_keeps_the_operator( + tmp_path, monkeypatch, druks_db +): with _header_client(tmp_path) as client: response = _connect( client, @@ -321,14 +331,14 @@ def test_reconnect_records_provider_email_but_keeps_the_operator(tmp_path, monke ) assert response.status_code == 200 assert response.json()["username"] == "me@example.com" - account = Account.get_for_username("me@example.com") - codex = HarnessConnection.get_for_account("codex", account.id) + account = await Account.get_for_username("me@example.com") + codex = await HarnessConnection.get_for_account("codex", account.id) assert codex.provider_email == "corp-seat@corp.com" -def test_connection_flow_rejects_a_bearer(tmp_path, druks_db): - agent = Account.get_or_create("agent@example.com") - _, token = PersonalAccessToken.create(account_id=agent.id, name="agent") +async def test_connection_flow_rejects_a_bearer(tmp_path, druks_db): + agent = await Account.get_or_create("agent@example.com") + _, token = await PersonalAccessToken.create(account_id=agent.id, name="agent") with _client(tmp_path) as client: response = client.post( "/api/harnesses/claude/connection/start", diff --git a/backend/tests/test_identity_jwt.py b/backend/tests/test_identity_jwt.py index 7eb20385..3626ad3e 100644 --- a/backend/tests/test_identity_jwt.py +++ b/backend/tests/test_identity_jwt.py @@ -62,14 +62,14 @@ def _jwt_client(tmp_path: Path) -> TestClient: return TestClient(app) -def test_a_valid_assertion_open_enrolls_its_subject(tmp_path, druks_db): +async def test_a_valid_assertion_open_enrolls_its_subject(tmp_path, druks_db): with _jwt_client(tmp_path) as client: response = client.get("/api/auth/me", headers={HEADER: _token()}) assert response.status_code == 200 assert response.json()["account"]["username"] == "op@example.com" other = client.get("/api/auth/me", headers={HEADER: _token(email="two@example.com")}) assert other.status_code == 200 - usernames = {account.username for account in Account.list_non_system()} + usernames = {account.username for account in await Account.list_non_system()} assert usernames == {"op@example.com", "two@example.com"} @@ -86,13 +86,13 @@ def test_a_valid_assertion_open_enrolls_its_subject(tmp_path, druks_db): "not.a.jwt", ], ) -def test_a_bad_assertion_rejects_without_enrolling(tmp_path, druks_db, token): +async def test_a_bad_assertion_rejects_without_enrolling(tmp_path, druks_db, token): with _jwt_client(tmp_path) as client: response = client.get("/api/auth/me", headers={HEADER: token}) assert response.status_code == 401 # Only the failure class reaches the caller — never token material. assert token.split(".")[1] not in response.json()["detail"] - assert not Account.list_non_system() + assert not await Account.list_non_system() def test_none_mode_multi_kid_document_serves_the_matching_key(tmp_path, druks_db): @@ -100,9 +100,9 @@ def test_none_mode_multi_kid_document_serves_the_matching_key(tmp_path, druks_db assert client.get("/api/auth/me", headers={HEADER: _token()}).status_code == 200 -def test_bearer_precedence_survives_jwt_mode(tmp_path, druks_db): - agent = Account.get_or_create("agent@example.com") - _, token = PersonalAccessToken.create(account_id=agent.id, name="agent") +async def test_bearer_precedence_survives_jwt_mode(tmp_path, druks_db): + agent = await Account.get_or_create("agent@example.com") + _, token = await PersonalAccessToken.create(account_id=agent.id, name="agent") with _jwt_client(tmp_path) as client: # A valid bearer wins over any assertion, even a garbage one. response = client.get( diff --git a/backend/tests/test_manifest.py b/backend/tests/test_manifest.py index 46f0407c..510de605 100644 --- a/backend/tests/test_manifest.py +++ b/backend/tests/test_manifest.py @@ -19,7 +19,7 @@ _TOKEN = "lin_secret_value" -def _build( +async def _build( *, harness: Harness | None = None, mcp_servers: tuple[McpServer, ...] = (), @@ -29,11 +29,11 @@ def _build( # get_manifest never touches the live sandbox, so the harness builds # without sandbox settings — the same shape argv unit tests use. harness = harness or ClaudeHarness(model="claude-opus-4-8", fast_mode=False, effort=None) - return harness.get_manifest(mcp_servers=mcp_servers, skills=skills, extra_env=extra_env) + return await harness.get_manifest(mcp_servers=mcp_servers, skills=skills, extra_env=extra_env) -def _seed_skills(*names: str, disabled: tuple[str, ...] = ()) -> None: - collection = SkillCollection.create( +async def _seed_skills(*names: str, disabled: tuple[str, ...] = ()) -> None: + collection = await SkillCollection.create( source="test", name="test skills", skills=[ @@ -46,11 +46,11 @@ def _seed_skills(*names: str, disabled: tuple[str, ...] = ()) -> None: skill.enabled = False -def test_manifest_records_the_delivered_capability_set(druks_db): +async def test_manifest_records_the_delivered_capability_set(druks_db): """Records model, harness, each MCP server's declared/delivered/token presence, and the delivered skills.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - _seed_skills("alpha", "beta") + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await _seed_skills("alpha", "beta") linear = McpServer(name="linear", url=_LINEAR_URL, bearer_token_env_var=_LINEAR_ENV) github = McpServer( @@ -60,7 +60,7 @@ def test_manifest_records_the_delivered_capability_set(druks_db): ) # Both servers delivered with their token — github is Ship's own # requirement (get_required_mcp_servers), so it reads delivered but not declared. - manifest = _build( + manifest = await _build( mcp_servers=(linear, github), skills=("alpha",), extra_env={ @@ -85,12 +85,12 @@ def test_manifest_records_the_delivered_capability_set(druks_db): assert github_entry["token_present"] is True -def test_missing_mcp_token_records_absence(druks_db): +async def test_missing_mcp_token_records_absence(druks_db): """A delivered server whose bearer var is absent from the run env reads token_present False — recorded, not failed.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - manifest = _build( + manifest = await _build( mcp_servers=(McpServer(name="linear", url=_LINEAR_URL, bearer_token_env_var=_LINEAR_ENV),), extra_env={}, ) @@ -100,13 +100,13 @@ def test_missing_mcp_token_records_absence(druks_db): assert linear_entry["token_present"] is False -def test_declared_but_undelivered_server_still_reads_declared(druks_db): +async def test_declared_but_undelivered_server_still_reads_declared(druks_db): """An enabled registry server is declared even on a call that didn't deliver it — declared True, delivered False, so the manifest shows exactly what this call ran without.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - manifest = _build() + manifest = await _build() linear_entry = next(s for s in manifest["mcp_servers"] if s["name"] == "linear") assert linear_entry["declared"] is True @@ -114,18 +114,18 @@ def test_declared_but_undelivered_server_still_reads_declared(druks_db): assert linear_entry["token_present"] is False -def test_records_the_delivered_server_not_the_registry_duplicate(druks_db): +async def test_records_the_delivered_server_not_the_registry_duplicate(druks_db): """When a workspace requires a server under an enabled entry's name, the workspace's wins delivery — the manifest records what the harness actually ran (the delivered url/env var), not the registry's values.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) required = McpServer( name="linear", url="https://required.internal/linear", bearer_token_env_var="REQUIRED_LINEAR_TOKEN", ) - manifest = _build(mcp_servers=(required,), extra_env={"REQUIRED_LINEAR_TOKEN": "s"}) + manifest = await _build(mcp_servers=(required,), extra_env={"REQUIRED_LINEAR_TOKEN": "s"}) linear_entry = next(s for s in manifest["mcp_servers"] if s["name"] == "linear") assert linear_entry["url"] == "https://required.internal/linear" @@ -135,55 +135,55 @@ def test_records_the_delivered_server_not_the_registry_duplicate(druks_db): assert linear_entry["token_present"] is True -def test_hash_is_stable_for_identical_capabilities(druks_db): +async def test_hash_is_stable_for_identical_capabilities(druks_db): """Identical capability sets hash the same; changing the model, MCP token availability, or the delivered skill set moves the hash.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) linear = McpServer(name="linear", url=_LINEAR_URL, bearer_token_env_var=_LINEAR_ENV) with_token = {"mcp_servers": (linear,), "extra_env": {_LINEAR_ENV: _TOKEN}} - baseline = _build(**with_token) - assert _build(**with_token)["manifest_hash"] == baseline["manifest_hash"] + baseline = await _build(**with_token) + assert (await _build(**with_token))["manifest_hash"] == baseline["manifest_hash"] - different_model = _build( + different_model = await _build( harness=ClaudeHarness(model="claude-sonnet-5", fast_mode=False, effort=None), **with_token, ) assert different_model["manifest_hash"] != baseline["manifest_hash"] - without_token = _build(mcp_servers=(linear,), extra_env={}) + without_token = await _build(mcp_servers=(linear,), extra_env={}) assert without_token["manifest_hash"] != baseline["manifest_hash"] -def test_curated_skills_change_the_recorded_set_and_the_hash(druks_db): +async def test_curated_skills_change_the_recorded_set_and_the_hash(druks_db): """Curated manifests omit disabled skills and hash each delivered set.""" - _seed_skills("alpha", "beta", disabled=("beta",)) - with_enabled = _build(skills=("alpha", "beta")) - disabled_only = _build(skills=("beta",)) + await _seed_skills("alpha", "beta", disabled=("beta",)) + with_enabled = await _build(skills=("alpha", "beta")) + disabled_only = await _build(skills=("beta",)) assert with_enabled["skills_delivered"] == ["alpha"] assert disabled_only["skills_delivered"] == [] assert with_enabled["manifest_hash"] != disabled_only["manifest_hash"] -def test_manifest_records_token_presence_never_the_value(druks_db): +async def test_manifest_records_token_presence_never_the_value(druks_db): """The secret token never lands in the manifest — only its env-var name and a presence boolean.""" - models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await models.McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) linear = McpServer(name="linear", url=_LINEAR_URL, bearer_token_env_var=_LINEAR_ENV) - manifest = _build(mcp_servers=(linear,), extra_env={_LINEAR_ENV: _TOKEN}) + manifest = await _build(mcp_servers=(linear,), extra_env={_LINEAR_ENV: _TOKEN}) serialized = json.dumps(manifest) assert _TOKEN not in serialized assert _LINEAR_ENV in serialized # the var name is safe to record -def test_manifest_stays_presence_only_for_a_declared_header_server(druks_db): +async def test_manifest_stays_presence_only_for_a_declared_header_server(druks_db): """A server delivered with declared headers records the same presence-only entry: no header value — plain or secret — lands in the manifest, and a bearer-less server simply reads token_present False.""" - models.McpServer.create( + await models.McpServer.create( name="grafana", url="https://mcp.grafana.com/mcp", token_source="", @@ -197,7 +197,7 @@ def test_manifest_stays_presence_only_for_a_declared_header_server(druks_db): env_headers={"X-Api-Key": "MCP_GRAFANA_HEADER_X_API_KEY"}, ) - manifest = _build( + manifest = await _build( mcp_servers=(delivered,), extra_env={"MCP_GRAFANA_HEADER_X_API_KEY": "grafana-api-secret"}, ) @@ -211,8 +211,8 @@ def test_manifest_stays_presence_only_for_a_declared_header_server(druks_db): assert grafana_entry["token_present"] is False -def test_persist_writes_manifest_into_the_call_dir(tmp_path, druks_db): - manifest = _build() +async def test_persist_writes_manifest_into_the_call_dir(tmp_path, druks_db): + manifest = await _build() path = persist_manifest(tmp_path, call_id="call-1", manifest=manifest) @@ -220,16 +220,16 @@ def test_persist_writes_manifest_into_the_call_dir(tmp_path, druks_db): assert json.loads(path.read_text())["manifest_hash"] == manifest["manifest_hash"] -def test_manifest_surfaces_in_agent_call_files(tmp_path, druks_db): +async def test_manifest_surfaces_in_agent_call_files(tmp_path, druks_db): """A written manifest.json is inventoried on the call's transcript files, in the manifest slot under its downloadable file name.""" - note = Note.create(body="manifest") - run = seed_run(druks_db, kind=Summarize.kind, subject=note) - call = seed_call(druks_db, run, "summarize", status="running") - manifest = _build() + note = await Note.create(body="manifest") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + call = await seed_call(druks_db, run, "summarize", status="running") + manifest = await _build() with mock.patch("druks.durable.models.load_settings", return_value=make_settings(tmp_path)): persist_manifest(call.call_dir.parent, call_id=call.call_dir.name, manifest=manifest) - files = get_agent_call_files(call.id) + files = await get_agent_call_files(call.id) assert files.manifest assert files.manifest.name == "manifest.json" diff --git a/backend/tests/test_mcp_endpoint.py b/backend/tests/test_mcp_endpoint.py index e8da82d6..e6f65c4d 100644 --- a/backend/tests/test_mcp_endpoint.py +++ b/backend/tests/test_mcp_endpoint.py @@ -13,10 +13,9 @@ from druks.durable.models import Artifact, Run from druks.mcp.exceptions import InvalidAgentToolError from druks.mcp.server import create_mcp_app -from druks.testing import configure_app_for_test, make_settings +from druks.testing import asgi_client, configure_app_for_test, make_settings from druks.usage.models import UsageScrape from fastapi import APIRouter, FastAPI -from fastapi.testclient import TestClient from fastmcp.client import Client from fastmcp.client.transports import StreamableHttpTransport from starlette.routing import Route @@ -70,13 +69,13 @@ async def live(app): @pytest.fixture -def account(druks_db): - return Account.get_or_create("op@example.com") +async def account(druks_db): + return await Account.get_or_create("op@example.com") @pytest.fixture -def pat_token(account): - _, token = PersonalAccessToken.create(account_id=account.id, name="agent") +async def pat_token(account): + _, token = await PersonalAccessToken.create(account_id=account.id, name="agent") return token @@ -119,7 +118,7 @@ def _wire_size(structured: dict) -> int: async def test_mcp_rejects_missing_and_dead_tokens(app, account, druks_db): - row, token = PersonalAccessToken.create(account_id=account.id, name="agent") + row, token = await PersonalAccessToken.create(account_id=account.id, name="agent") async with httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://druks.test" ) as wire: @@ -137,12 +136,12 @@ async def test_mcp_rejects_missing_and_dead_tokens(app, account, druks_db): bearer = {**_WIRE_HEADERS, "Authorization": f"Bearer {token}"} row.expires_at = datetime.now(UTC) - timedelta(days=1) - druks_db.flush() + await druks_db.flush() expired = await wire.post("/mcp", json=_INIT, headers=bearer) assert expired.status_code == 401 row.expires_at = datetime.now(UTC) + timedelta(days=1) - row.revoke() + await row.revoke() revoked = await wire.post("/mcp", json=_INIT, headers=bearer) assert revoked.status_code == 401 @@ -327,10 +326,10 @@ async def test_app_agent_route_derives_the_namespaced_tool( async def test_claims_resolve_the_calling_account(app, druks_db): # get_usage must answer as the token's account — the forwarded bearer. - mine = Account.get_or_create("op@example.com") - theirs = Account.get_or_create("peer@example.com") - _, my_token = PersonalAccessToken.create(account_id=mine.id, name="mine") - _, their_token = PersonalAccessToken.create(account_id=theirs.id, name="theirs") + mine = await Account.get_or_create("op@example.com") + theirs = await Account.get_or_create("peer@example.com") + _, my_token = await PersonalAccessToken.create(account_id=mine.id, name="mine") + _, their_token = await PersonalAccessToken.create(account_id=theirs.id, name="theirs") druks_db.add( UsageScrape( harness="codex", @@ -339,7 +338,7 @@ async def test_claims_resolve_the_calling_account(app, druks_db): five_hour_percent_left=42, ) ) - druks_db.flush() + await druks_db.flush() async with live(app), _client(app, my_token) as client: usage = (await client.call_tool("get_usage", {})).structured_content @@ -355,8 +354,8 @@ async def test_claims_resolve_the_calling_account(app, druks_db): async def test_gate_cycle_reads_answers_and_reports_stale_rounds( app, pat_token, druks_db, resume_spy ): - item = make_test_note() - run = seed_note_run( + item = await make_test_note() + run = await seed_note_run( druks_db, note=item, state="parked", @@ -364,7 +363,7 @@ async def test_gate_cycle_reads_answers_and_reports_stale_rounds( input_request=dict(_IN_APP_ASK), ) run.input_requested_at = datetime.now(UTC) - druks_db.flush() + await druks_db.flush() async with live(app), _client(app, pat_token) as client: gate = (await client.call_tool("get_gate", {"run": run.id})).structured_content @@ -389,7 +388,7 @@ async def test_gate_cycle_reads_answers_and_reports_stale_rounds( assert resume_spy == [{"id": run.id, "action": "approve", "answers": {}, "note": ""}] run.answer_parked_at = run.input_requested_at - druks_db.flush() + await druks_db.flush() repeat = ( await client.call_tool( "answer_gate", @@ -405,8 +404,8 @@ async def test_gate_cycle_reads_answers_and_reports_stale_rounds( ) assert missing["code"] == "RUN_NOT_FOUND" - external_item = make_test_note() - external = seed_note_run( + external_item = await make_test_note() + external = await seed_note_run( druks_db, note=external_item, state="parked", @@ -414,19 +413,19 @@ async def test_gate_cycle_reads_answers_and_reports_stale_rounds( input_request={"presentation": "external"}, ) external.input_requested_at = datetime.now(UTC) - druks_db.flush() + await druks_db.flush() unanswerable = await _call_error(client, "get_gate", {"run": external.id}) assert unanswerable["code"] == "GATE_NOT_ANSWERABLE" async def test_get_agent_call_serves_bounded_tails(app, pat_token, druks_db): - call = seed_note_agent_run() + call = await seed_note_agent_run() call_dir = call.call_dir call_dir.mkdir(parents=True, exist_ok=True) (call_dir / "stdout.jsonl").write_bytes(b"s" * 20480) (call_dir / "stderr.log").write_bytes(b"e" * 10240) - finish_agent_run(call, last_error="boom " * 100) - Artifact.record( + await finish_agent_run(call, last_error="boom " * 100) + await Artifact.record( call_dir=call_dir, call_id=call.id, kind="markdown", title="Out", content="a" * 10240 ) @@ -453,12 +452,12 @@ async def _spy(self, *, failure): cancels.append({"id": self.id, "failure": failure}) monkeypatch.setattr(Run, "cancel", _spy) - item = make_test_note() - active = seed_note_run(druks_db, note=item, state="running") - done_item = make_test_note() - done = seed_note_run(druks_db, note=done_item, state="finished") - gone_item = make_test_note() - gone = seed_note_run(druks_db, note=gone_item, state="cancelled") + item = await make_test_note() + active = await seed_note_run(druks_db, note=item, state="running") + done_item = await make_test_note() + done = await seed_note_run(druks_db, note=done_item, state="finished") + gone_item = await make_test_note() + gone = await seed_note_run(druks_db, note=gone_item, state="cancelled") async with live(app), _client(app, pat_token) as client: cancelled = ( @@ -497,9 +496,10 @@ async def set_status(self, key, status): async def aclose(self): pass - monkeypatch.setattr( - Ship, "get_tracker", classmethod(lambda cls, source=None: _UnknownTicketTracker()) - ) + async def _get_tracker(cls, source=None): + return _UnknownTicketTracker() + + monkeypatch.setattr(Ship, "get_tracker", classmethod(_get_tracker)) async with live(app), _client(app, pat_token) as client: error = await _call_error(client, "ship_start", {"ticket": "ENG-9999"}) @@ -519,7 +519,7 @@ async def test_get_usage_reads_within_budget(app, pat_token): assert _wire_size(usage) <= 4 * 1024 -def test_lifespan_composes_the_endpoint_once(app, monkeypatch): +async def test_lifespan_composes_the_endpoint_once(app, monkeypatch): entered = [] original = mcp_app.router.lifespan_context @@ -530,16 +530,16 @@ async def counting(scope_app): yield monkeypatch.setattr(mcp_app.router, "lifespan_context", counting) - with TestClient(app) as client: - assert client.get("/api/system/health").status_code == 200 + async with app.router.lifespan_context(app), asgi_client(app) as client: + assert (await client.get("/api/system/health")).status_code == 200 assert entered == [1] -def test_mcp_server_registry_routes_stay_untouched(tmp_path, druks_db, monkeypatch): +async def test_mcp_server_registry_routes_stay_untouched(tmp_path, druks_db, monkeypatch): monkeypatch.setenv("DRUKS_DATA_DIR", str(tmp_path)) app = configure_app_for_test(settings=make_settings(tmp_path)) - with TestClient(app) as client: - listed = client.get("/api/mcp-servers") + async with asgi_client(app) as client: + listed = await client.get("/api/mcp-servers") assert listed.status_code == 200 # The inbound endpoint never joins the outbound server registry. assert "druks" not in {server["name"] for server in listed.json()} diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py index 25b0dd90..93f933f2 100644 --- a/backend/tests/test_mcp_oauth.py +++ b/backend/tests/test_mcp_oauth.py @@ -118,27 +118,27 @@ def _register_oauth_server(name: str = _NAME, enabled: bool = True) -> None: ) -def _store_grant( +async def _store_grant( refresh_token: str = "rt-1", *, account_id: str = SYSTEM_ACCOUNT_ID, identity_mode: IdentityMode = IdentityMode.SHARED, ) -> OauthConnection: - server = McpServer.get_for_name(_NAME) + server = await McpServer.get_for_name(_NAME) if not server: - server = McpServer.create( + server = await McpServer.create( name=_NAME, url=_SERVER_URL, token_source=TokenSource.OAUTH, ) server.identity_mode = identity_mode - McpClientRegistration.store( + await McpClientRegistration.store( server_id=server.id, account_id=account_id, token_endpoint=f"{_AUTH_BASE}/token", client_id="client-123", ) - return OauthConnection.create( + return await OauthConnection.create( provider=f"mcp:{_NAME}", account_id=account_id, refresh_token=refresh_token, @@ -150,12 +150,12 @@ def _state_key(state: str) -> str: return f"oauth:connect:{state}" -def _token_key(account_id: str) -> str: - return f"mcp:{_NAME}:access_token:{oauth.get_connection(_NAME, account_id).id}" +async def _token_key(account_id: str) -> str: + return f"mcp:{_NAME}:access_token:{(await oauth.get_connection(_NAME, account_id)).id}" -def _lock_key(account_id: str) -> str: - return f"mcp:{_NAME}:refresh_lock:{oauth.get_connection(_NAME, account_id).id}" +async def _lock_key(account_id: str) -> str: + return f"mcp:{_NAME}:refresh_lock:{(await oauth.get_connection(_NAME, account_id)).id}" @pytest.mark.parametrize( @@ -310,10 +310,10 @@ async def test_complete_connect_exchanges_code_and_stores_the_grant(auth_server, name = await oauth.complete_connect(state=state, code="code-1") assert name == _NAME - grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + grant = await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert grant.refresh_token.decrypt() == "rt-1" assert grant.identity == {"email": "op@linear.test"} - registration = McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + registration = await McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) assert registration.client_id == "client-123" assert registration.token_endpoint == f"{_AUTH_BASE}/token" exchange = auth_server.token_requests[0] @@ -324,7 +324,7 @@ async def test_complete_connect_exchanges_code_and_stores_the_grant(auth_server, # Nothing is cached at connect (the grant is real only once this commits); # the first delivery mints from it, carrying the grant's resource binding. - assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) + assert not await get_client().get(await _token_key(SYSTEM_ACCOUNT_ID)) assert await oauth.get_access_token(_NAME, SYSTEM_ACCOUNT_ID) == "at-1" refresh = auth_server.token_requests[1] assert refresh["grant_type"] == "refresh_token" @@ -349,7 +349,7 @@ async def test_off_issuer_userinfo_is_dropped_and_logged(auth_server, druks_db, with caplog.at_level("WARNING"): await oauth.complete_connect(state=state, code="code-1") - assert oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID).identity == {} + assert (await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID)).identity == {} assert "evil.test" in caplog.text @@ -366,12 +366,12 @@ async def test_complete_connect_without_refresh_token_stores_nothing(auth_server with pytest.raises(OauthConnectError, match="no refresh token"): await oauth.complete_connect(state=state, code="code-1") - assert not oauth.list_connections(_NAME) + assert not await oauth.list_connections(_NAME) async def test_reconsent_replaces_the_grant_and_evicts_the_stale_token(auth_server, druks_db): - _store_grant(refresh_token="rt-stale") - await get_client().set(_token_key(SYSTEM_ACCOUNT_ID), "at-stale") + await _store_grant(refresh_token="rt-stale") + await get_client().set(await _token_key(SYSTEM_ACCOUNT_ID), "at-stale") url = await oauth.begin_connect( _NAME, @@ -383,17 +383,17 @@ async def test_reconsent_replaces_the_grant_and_evicts_the_stale_token(auth_serv state = dict(parse_qsl(urlparse(url).query))["state"] await oauth.complete_connect(state=state, code="code-1") - grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + grant = await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert grant.refresh_token.decrypt() == "rt-1" # The stale narrow token must not keep serving until its TTL runs out. - assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) + assert not await get_client().get(await _token_key(SYSTEM_ACCOUNT_ID)) async def test_reconnect_after_disconnect_creates_a_new_grant(auth_server, druks_db): - _store_grant(refresh_token="rt-stale") - revoked = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + await _store_grant(refresh_token="rt-stale") + revoked = await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) await oauth.disconnect(_NAME, SYSTEM_ACCOUNT_ID) - assert not oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + assert not await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) url = await oauth.begin_connect( _NAME, @@ -407,18 +407,18 @@ async def test_reconnect_after_disconnect_creates_a_new_grant(auth_server, druks from druks.database import db_session - db_session().expire_all() + db_session().expunge_all() # A re-connect creates a new grant. The revoked row stays as history, # so at most one live connection holds the (server, account) slot. - grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + grant = await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert grant.id != revoked.id assert grant.refresh_token.decrypt() == "rt-1" assert revoked.revoked_at async def test_two_shared_connects_converge_on_one_grant(auth_server, druks_db): - first = Account.get_or_create("first@example.com") - second = Account.get_or_create("second@example.com") + first = await Account.get_or_create("first@example.com") + second = await Account.get_or_create("second@example.com") for account in (first, second): url = await oauth.begin_connect( @@ -431,14 +431,14 @@ async def test_two_shared_connects_converge_on_one_grant(auth_server, druks_db): state = dict(parse_qsl(urlparse(url).query))["state"] await oauth.complete_connect(state=state, code=account.id) - grants = oauth.list_connections(_NAME) - assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.SHARED + grants = await oauth.list_connections(_NAME) + assert (await McpServer.get_for_name(_NAME)).identity_mode == IdentityMode.SHARED assert [grant.account_id for grant in grants] == [SYSTEM_ACCOUNT_ID] async def test_two_per_user_connects_store_two_grants(auth_server, druks_db): - first = Account.get_or_create("first@example.com") - second = Account.get_or_create("second@example.com") + first = await Account.get_or_create("first@example.com") + second = await Account.get_or_create("second@example.com") for account in (first, second): url = await oauth.begin_connect( @@ -451,14 +451,14 @@ async def test_two_per_user_connects_store_two_grants(auth_server, druks_db): state = dict(parse_qsl(urlparse(url).query))["state"] await oauth.complete_connect(state=state, code=account.id) - grants = oauth.list_connections(_NAME) - assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.PER_USER + grants = await oauth.list_connections(_NAME) + assert (await McpServer.get_for_name(_NAME)).identity_mode == IdentityMode.PER_USER assert {grant.account_id for grant in grants} == {first.id, second.id} async def test_a_later_connect_stores_under_the_claimed_mode(auth_server, druks_db): - first = Account.get_or_create("first@example.com") - second = Account.get_or_create("second@example.com") + first = await Account.get_or_create("first@example.com") + second = await Account.get_or_create("second@example.com") first_url = await oauth.begin_connect( _NAME, _SERVER_URL, @@ -479,8 +479,8 @@ async def test_a_later_connect_stores_under_the_claimed_mode(auth_server, druks_ await oauth.complete_connect(state=first_state, code="first") await oauth.complete_connect(state=second_state, code="second") - assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.SHARED - grant_accounts = {grant.account_id for grant in oauth.list_connections(_NAME)} + assert (await McpServer.get_for_name(_NAME)).identity_mode == IdentityMode.SHARED + grant_accounts = {grant.account_id for grant in await oauth.list_connections(_NAME)} assert grant_accounts == {SYSTEM_ACCOUNT_ID} @@ -488,7 +488,7 @@ async def test_a_later_connect_stores_under_the_claimed_mode(auth_server, druks_ async def test_get_refreshes_on_cache_miss_and_persists_rotation(auth_server, druks_db): - _store_grant(refresh_token="rt-old") + await _store_grant(refresh_token="rt-old") auth_server.token_response = { "access_token": "at-2", "refresh_token": "rt-new", @@ -504,7 +504,7 @@ async def test_get_refreshes_on_cache_miss_and_persists_rotation(auth_server, dr assert refresh["refresh_token"] == "rt-old" assert refresh["resource"] == _SERVER_URL # Rotation: the provider's new refresh token replaced the stored one. - stored = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + stored = await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert stored.refresh_token.decrypt() == "rt-new" # A second call within the TTL reuses the cache — no second refresh. @@ -518,16 +518,16 @@ async def test_get_without_grant_fails_loudly(druks_db): async def test_get_refresh_rejection_fails_loudly_and_evicts_the_cache(auth_server, druks_db): - _store_grant() + await _store_grant() auth_server.token_status = 400 with pytest.raises(GrantRefreshError, match=_NAME): await oauth.get_access_token(_NAME, SYSTEM_ACCOUNT_ID) - assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) + assert not await get_client().get(await _token_key(SYSTEM_ACCOUNT_ID)) async def test_get_rejects_a_malformed_token_response(auth_server, druks_db): - _store_grant() + await _store_grant() auth_server.token_malformed = True with pytest.raises(GrantRefreshError, match="malformed JSON"): @@ -535,7 +535,7 @@ async def test_get_rejects_a_malformed_token_response(auth_server, druks_db): async def test_get_rejects_a_token_response_without_an_access_token(auth_server, druks_db): - _store_grant() + await _store_grant() auth_server.token_response = {"refresh_token": "rt-2", "expires_in": 3600} with pytest.raises(GrantRefreshError, match="no access token"): @@ -548,14 +548,14 @@ async def test_get_losing_the_refresh_lock_polls_for_the_winners_token( # A second caller never refreshes (one rotation spender per server) and # never blocks the event loop: it polls until the winner's token appears # in the cache. The winner here is a concurrent task holding the lock. - _store_grant() + await _store_grant() redis = get_client() monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_INTERVAL_SECONDS", 0) - await redis.set(_lock_key(SYSTEM_ACCOUNT_ID), "1") + await redis.set(await _lock_key(SYSTEM_ACCOUNT_ID), "1") async def _winner_finishes(): - await redis.set(_token_key(SYSTEM_ACCOUNT_ID), "at-winner") - await redis.delete(_lock_key(SYSTEM_ACCOUNT_ID)) + await redis.set(await _token_key(SYSTEM_ACCOUNT_ID), "at-winner") + await redis.delete(await _lock_key(SYSTEM_ACCOUNT_ID)) winner = asyncio.create_task(_winner_finishes()) assert await oauth.get_access_token(_NAME, SYSTEM_ACCOUNT_ID) == "at-winner" @@ -564,26 +564,26 @@ async def _winner_finishes(): async def test_get_times_out_loudly_when_the_refresh_lock_never_frees(druks_db, monkeypatch): - _store_grant() + await _store_grant() monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_INTERVAL_SECONDS", 0) monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_ATTEMPTS", 3) - await get_client().set(_lock_key(SYSTEM_ACCOUNT_ID), "1") + await get_client().set(await _lock_key(SYSTEM_ACCOUNT_ID), "1") with pytest.raises(GrantRefreshError, match="concurrent refresh"): await oauth.get_access_token(_NAME, SYSTEM_ACCOUNT_ID) async def test_get_cache_and_refresh_lock_are_per_account(auth_server, druks_db): - first = Account.get_or_create("first@example.com") - second = Account.get_or_create("second@example.com") - _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) - _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) + first = await Account.get_or_create("first@example.com") + second = await Account.get_or_create("second@example.com") + await _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) + await _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - await redis.set(_lock_key(first.id), "1") + await redis.set(await _lock_key(first.id), "1") assert await oauth.get_access_token(_NAME, second.id) == "at-1" - assert not await redis.get(_token_key(first.id)) - assert await redis.get(_token_key(second.id)) == b"at-1" + assert not await redis.get(await _token_key(first.id)) + assert await redis.get(await _token_key(second.id)) == b"at-1" # --- delivery: the oauth branch of the fold --------------------------------- @@ -591,7 +591,7 @@ async def test_get_cache_and_refresh_lock_are_per_account(auth_server, druks_db) async def test_delivery_mints_and_injects_the_oauth_token(registry_state, auth_server, druks_db): _register_oauth_server() - _store_grant() + await _store_grant() kwargs = await Workspace(sandbox=_FakeSandbox()).with_mcp_servers( # type: ignore[arg-type] SYSTEM_ACCOUNT_ID @@ -609,7 +609,7 @@ async def test_delivery_fails_loudly_for_an_unconnected_enabled_oauth_server( registry_state, druks_db ): _register_oauth_server() - server = McpServer.create(name=_NAME, url=_SERVER_URL, token_source=TokenSource.OAUTH) + server = await McpServer.create(name=_NAME, url=_SERVER_URL, token_source=TokenSource.OAUTH) server.identity_mode = IdentityMode.SHARED with pytest.raises(MissingGrantError, match=_NAME): @@ -619,8 +619,8 @@ async def test_delivery_fails_loudly_for_an_unconnected_enabled_oauth_server( async def test_delivery_names_the_account_missing_its_per_user_grant(druks_db): - account = Account.get_or_create("run@example.com") - server = McpServer.create(name=_NAME, url=_SERVER_URL, token_source=TokenSource.OAUTH) + account = await Account.get_or_create("run@example.com") + server = await McpServer.create(name=_NAME, url=_SERVER_URL, token_source=TokenSource.OAUTH) server.identity_mode = IdentityMode.PER_USER with pytest.raises(MissingGrantError) as error: @@ -635,9 +635,9 @@ async def test_delivery_names_the_account_missing_its_per_user_grant(druks_db): async def test_delivery_without_a_run_account_uses_the_fallback(auth_server, druks_db): - fallback = Account.get_or_create("fallback@example.com") - UserSettings.get().set_fallback_account(fallback.id) - _store_grant(account_id=fallback.id, identity_mode=IdentityMode.PER_USER) + fallback = await Account.get_or_create("fallback@example.com") + await (await UserSettings.get()).set_fallback_account(fallback.id) + await _store_grant(account_id=fallback.id, identity_mode=IdentityMode.PER_USER) kwargs = await Workspace(sandbox=_FakeSandbox()).with_mcp_servers( # type: ignore[arg-type] None @@ -647,10 +647,10 @@ async def test_delivery_without_a_run_account_uses_the_fallback(auth_server, dru async def test_delivery_with_a_named_account_does_not_use_the_fallback(auth_server, druks_db): - fallback = Account.get_or_create("fallback@example.com") - named = Account.get_or_create("named@example.com") - UserSettings.get().set_fallback_account(fallback.id) - _store_grant(account_id=fallback.id, identity_mode=IdentityMode.PER_USER) + fallback = await Account.get_or_create("fallback@example.com") + named = await Account.get_or_create("named@example.com") + await (await UserSettings.get()).set_fallback_account(fallback.id) + await _store_grant(account_id=fallback.id, identity_mode=IdentityMode.PER_USER) with pytest.raises(MissingGrantError) as error: await Workspace(sandbox=_FakeSandbox()).with_mcp_servers( # type: ignore[arg-type] @@ -697,7 +697,9 @@ def test_connect_route_returns_the_consent_url(tmp_path, registry_state, auth_se assert response.json()["authorizationUrl"].startswith(f"{_AUTH_BASE}/authorize?") -def test_callback_route_completes_the_connect(tmp_path, registry_state, auth_server, druks_db): +async def test_callback_route_completes_the_connect( + tmp_path, registry_state, auth_server, druks_db +): _register_oauth_server(enabled=False) settings = make_settings(tmp_path, urls={"endpoint": _ENDPOINT}) @@ -713,9 +715,9 @@ def test_callback_route_completes_the_connect(tmp_path, registry_state, auth_ser # The page notifies the opener tab, then closes itself. assert "BroadcastChannel('druks-mcp-connect')" in page.text assert "window.close()" in page.text - assert oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + assert await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) # Connecting is the explicit "use this server" — it enables too. - assert McpServer.get_for_name(_NAME).is_enabled is True + assert (await McpServer.get_for_name(_NAME)).is_enabled is True # Consent denied / unknown state both land loudly, storing nothing. assert ( @@ -737,20 +739,20 @@ async def test_disconnect_route_drops_grant_and_cache( tmp_path, registry_state, auth_server, druks_db ): _register_oauth_server() - _store_grant() + await _store_grant() await oauth.get_access_token(_NAME, SYSTEM_ACCOUNT_ID) - token_key = _token_key(SYSTEM_ACCOUNT_ID) + token_key = await _token_key(SYSTEM_ACCOUNT_ID) # The oauth engine's Redis client is bound to this test's loop; close it so the # route dials its own — the cached token lives in Redis either way. await close_client() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: assert client.delete(f"/api/mcp-servers/{_NAME}/grant").status_code == 204 - assert not oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) - assert not McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + assert not await oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + assert not await McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) # The mirror of connect-enables: no grant, no calls, so no dead entry # riding into VMs. - assert McpServer.get_for_name(_NAME).is_enabled is False + assert (await McpServer.get_for_name(_NAME)).is_enabled is False assert client.delete(f"/api/mcp-servers/{_NAME}/grant").status_code == 404 # Read the eviction on this test's own loop. Abandon whatever client the @@ -761,11 +763,11 @@ async def test_disconnect_route_drops_grant_and_cache( assert not await get_client().get(token_key) -def test_shared_disconnect_allows_per_user_reconnect( +async def test_shared_disconnect_allows_per_user_reconnect( tmp_path, registry_state, auth_server, druks_db ): _register_oauth_server() - operator = Account.get_or_create("op@example.com") + operator = await Account.get_or_create("op@example.com") settings = make_settings(tmp_path, urls={"endpoint": _ENDPOINT}) with TestClient(configure_app_for_test(settings=settings)) as client: @@ -796,26 +798,28 @@ def test_shared_disconnect_allows_per_user_reconnect( == 200 ) - assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.PER_USER - assert {grant.account_id for grant in oauth.list_connections(_NAME)} == {operator.id} + assert (await McpServer.get_for_name(_NAME)).identity_mode == IdentityMode.PER_USER + assert {grant.account_id for grant in await oauth.list_connections(_NAME)} == {operator.id} -def test_api_has_token_reflects_the_grant_and_leaks_no_secret(tmp_path, registry_state, druks_db): +async def test_api_has_token_reflects_the_grant_and_leaks_no_secret( + tmp_path, registry_state, druks_db +): _register_oauth_server() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: server = next(s for s in client.get("/api/mcp-servers").json() if s["name"] == _NAME) assert server["tokenSource"] == "oauth" assert server["hasToken"] is False - _store_grant(refresh_token="rt-secret-value") + await _store_grant(refresh_token="rt-secret-value") listed = client.get("/api/mcp-servers") server = next(s for s in listed.json() if s["name"] == _NAME) assert server["hasToken"] is True assert "rt-secret-value" not in listed.text -def test_connect_route_rejects_a_conflicting_identity_mode(tmp_path, druks_db): - _store_grant() +async def test_connect_route_rejects_a_conflicting_identity_mode(tmp_path, druks_db): + await _store_grant() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: response = client.post( @@ -827,10 +831,10 @@ def test_connect_route_rejects_a_conflicting_identity_mode(tmp_path, druks_db): assert IdentityMode.SHARED in response.json()["detail"] -def test_api_has_token_is_scoped_to_the_requesting_account(tmp_path, druks_db): - connected = Account.get_or_create("connected@example.com") - unconnected = Account.get_or_create("unconnected@example.com") - _store_grant(account_id=connected.id, identity_mode=IdentityMode.PER_USER) +async def test_api_has_token_is_scoped_to_the_requesting_account(tmp_path, druks_db): + connected = await Account.get_or_create("connected@example.com") + unconnected = await Account.get_or_create("unconnected@example.com") + await _store_grant(account_id=connected.id, identity_mode=IdentityMode.PER_USER) settings = make_settings( tmp_path, identity={"mode": "header", "header": "X-Test-User"}, @@ -857,13 +861,13 @@ def test_api_has_token_is_scoped_to_the_requesting_account(tmp_path, druks_db): async def test_per_user_disconnect_preserves_other_accounts_grant_and_cache(tmp_path, druks_db): - disconnected = Account.get_or_create("disconnect@example.com") - connected = Account.get_or_create("connected@example.com") - _store_grant(account_id=disconnected.id, identity_mode=IdentityMode.PER_USER) - _store_grant(account_id=connected.id, identity_mode=IdentityMode.PER_USER) + disconnected = await Account.get_or_create("disconnect@example.com") + connected = await Account.get_or_create("connected@example.com") + await _store_grant(account_id=disconnected.id, identity_mode=IdentityMode.PER_USER) + await _store_grant(account_id=connected.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - disconnected_key = _token_key(disconnected.id) - connected_key = _token_key(connected.id) + disconnected_key = await _token_key(disconnected.id) + connected_key = await _token_key(connected.id) await redis.set(disconnected_key, "disconnect-token") await redis.set(connected_key, "connected-token") await close_client() @@ -879,22 +883,22 @@ async def test_per_user_disconnect_preserves_other_accounts_grant_and_cache(tmp_ ) assert response.status_code == 204 - assert not oauth.get_connection(_NAME, disconnected.id) - assert oauth.get_connection(_NAME, connected.id) - assert McpServer.get_for_name(_NAME).is_enabled is True + assert not await oauth.get_connection(_NAME, disconnected.id) + assert await oauth.get_connection(_NAME, connected.id) + assert (await McpServer.get_for_name(_NAME)).is_enabled is True druks.redis._client = None assert not await get_client().get(disconnected_key) assert await get_client().get(connected_key) == b"connected-token" async def test_removal_drops_every_grant_and_cached_token(tmp_path, druks_db): - first = Account.get_or_create("first@example.com") - second = Account.get_or_create("second@example.com") - _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) - _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) + first = await Account.get_or_create("first@example.com") + second = await Account.get_or_create("second@example.com") + await _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) + await _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - first_key = _token_key(first.id) - second_key = _token_key(second.id) + first_key = await _token_key(first.id) + second_key = await _token_key(second.id) await redis.set(first_key, "first-token") await redis.set(second_key, "second-token") await close_client() @@ -902,8 +906,8 @@ async def test_removal_drops_every_grant_and_cached_token(tmp_path, druks_db): with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: assert client.delete(f"/api/mcp-servers/{_NAME}").status_code == 204 - assert not McpServer.get_for_name(_NAME) - assert not oauth.list_connections(_NAME) + assert not await McpServer.get_for_name(_NAME) + assert not await oauth.list_connections(_NAME) druks.redis._client = None assert not await get_client().get(first_key) assert not await get_client().get(second_key) diff --git a/backend/tests/test_mcp_registry.py b/backend/tests/test_mcp_registry.py index 12c3e34c..9297c025 100644 --- a/backend/tests/test_mcp_registry.py +++ b/backend/tests/test_mcp_registry.py @@ -312,7 +312,9 @@ def test_registry_search_route_maps_unavailability_to_502(tmp_path, monkeypatch, assert "registry search" in response.json()["detail"] -def test_add_from_registry_writes_the_row_and_redacts_the_secret(tmp_path, monkeypatch, druks_db): +async def test_add_from_registry_writes_the_row_and_redacts_the_secret( + tmp_path, monkeypatch, druks_db +): with _client_with_registry(tmp_path, monkeypatch, _ACME_ENTRY) as client: created = client.post( "/api/mcp-servers/registry", @@ -337,14 +339,16 @@ def test_add_from_registry_writes_the_row_and_redacts_the_secret(tmp_path, monke # The row: url from the registry (never the client), values split by the # spec's secrecy — the plain one readable, the secret one ciphertext at # rest and redacted in repr. - row = McpServer.get_for_name("observer") + row = await McpServer.get_for_name("observer") assert row.url == "https://mcp.acme.com/mcp" assert row.headers == {"X-Region": "eu"} assert "acme-api-secret" not in repr(row.secret_headers) assert row.secret_headers["X-Api-Key"] == "acme-api-secret" -def test_add_from_registry_oauth_candidate_ships_dark_and_connects(tmp_path, monkeypatch, druks_db): +async def test_add_from_registry_oauth_candidate_ships_dark_and_connects( + tmp_path, monkeypatch, druks_db +): with _client_with_registry(tmp_path, monkeypatch, _GRAFANA) as client: created = client.post( "/api/mcp-servers/registry", @@ -389,11 +393,11 @@ async def fake_begin_connect(name, server_url, endpoint, *, account_id, identity assert begun[0][3] assert begun[0][4] == IdentityMode.PER_USER - row = McpServer.get_for_name("grafana") + row = await McpServer.get_for_name("grafana") assert row.headers == {"X-Grafana-URL": "https://acme.grafana.net"} -def test_add_from_registry_rejects_missing_required_and_unknown_headers( +async def test_add_from_registry_rejects_missing_required_and_unknown_headers( tmp_path, monkeypatch, druks_db ): with _client_with_registry(tmp_path, monkeypatch, _ACME_ENTRY) as client: @@ -422,7 +426,7 @@ def test_add_from_registry_rejects_missing_required_and_unknown_headers( assert unknown.status_code == 422 assert "X-Bogus" in unknown.json()["detail"] - assert not McpServer.get_for_name("observer") + assert not await McpServer.get_for_name("observer") def test_add_from_registry_rejects_an_entry_without_an_http_remote(tmp_path, monkeypatch, druks_db): @@ -437,18 +441,18 @@ def test_add_from_registry_rejects_an_entry_without_an_http_remote(tmp_path, mon assert "not installable" in created.json()["detail"] -def test_removing_a_connected_row_drops_its_grant(tmp_path, monkeypatch, druks_db): +async def test_removing_a_connected_row_drops_its_grant(tmp_path, monkeypatch, druks_db): with _client_with_registry(tmp_path, monkeypatch, _GRAFANA) as client: client.post( "/api/mcp-servers/registry", json={"name": "grafana", "registry": "io.github.grafana/mcp-grafana", "headers": {}}, ) - OauthConnection.create( + await OauthConnection.create( provider="mcp:grafana", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt", scopes=[] ) assert client.delete("/api/mcp-servers/grafana").status_code == 204 # An orphan grant would revive as this name's credential on re-add. - assert not McpServer.get_for_name("grafana") - assert not OauthConnection.list_for_provider("mcp:grafana") + assert not await McpServer.get_for_name("grafana") + assert not await OauthConnection.list_for_provider("mcp:grafana") diff --git a/backend/tests/test_mcp_servers.py b/backend/tests/test_mcp_servers.py index 5251cb81..cecc529b 100644 --- a/backend/tests/test_mcp_servers.py +++ b/backend/tests/test_mcp_servers.py @@ -17,9 +17,8 @@ from druks.mcp.models import McpServer from druks.sandbox.datastructures import RequiredMcpServer from druks.settings import PACKAGED_MCP_CATALOG -from druks.testing import configure_app_for_test, make_settings +from druks.testing import asgi_client, configure_app_for_test, make_settings from druks.workspaces import Workspace -from fastapi.testclient import TestClient _LINEAR_URL = "https://mcp.linear.app/mcp" _TOKEN = "lin_secret_value" @@ -61,44 +60,44 @@ def get_required_mcp_servers(self) -> tuple[RequiredMcpServer, ...]: # --- custom servers: CRUD + enable/disable ------------------------------- -def test_create_lists_and_deletes(druks_db): - server = McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) +async def test_create_lists_and_deletes(druks_db): + server = await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - by_name = McpServer.get_for_name("linear") + by_name = await McpServer.get_for_name("linear") assert by_name assert by_name.id == server.id - assert "linear" in {s.name for s in McpServer.list_all()} + assert "linear" in {s.name for s in await McpServer.list_all()} - server.delete() - assert not McpServer.get_for_name("linear") + await server.delete() + assert not await McpServer.get_for_name("linear") -def test_enable_disable_moves_in_and_out_of_the_enabled_set(druks_db): - server = McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - assert "linear" in {s["name"] for s in McpServer.list_enabled()} +async def test_enable_disable_moves_in_and_out_of_the_enabled_set(druks_db): + server = await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + assert "linear" in {s["name"] for s in await McpServer.list_enabled()} server.is_enabled = False - druks_db.flush() - assert "linear" not in {s["name"] for s in McpServer.list_enabled()} + await druks_db.flush() + assert "linear" not in {s["name"] for s in await McpServer.list_enabled()} server.is_enabled = True - druks_db.flush() - assert "linear" in {s["name"] for s in McpServer.list_enabled()} + await druks_db.flush() + assert "linear" in {s["name"] for s in await McpServer.list_enabled()} # --- name validity: one identifier, shell/TOML-safe ---------------------- -def test_create_rejects_names_that_break_env_or_config(druks_db): +async def test_create_rejects_names_that_break_env_or_config(druks_db): # A hyphen breaks the sourced ``KEY='value'`` env line and the codex TOML key # path; a leading digit and uppercase are rejected for the same reason. for bad in ("linear-app", "1linear", "Linear", "linear.app", "linear app"): with pytest.raises(InvalidServerNameError, match="Invalid MCP server name"): - McpServer.create(name=bad, url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name=bad, url=_LINEAR_URL, token=_TOKEN) -def test_valid_name_derives_shell_safe_env_var(druks_db): - server = McpServer.create(name="linear_app", url=_LINEAR_URL, token=_TOKEN) +async def test_valid_name_derives_shell_safe_env_var(druks_db): + server = await McpServer.create(name="linear_app", url=_LINEAR_URL, token=_TOKEN) # Every char of the derived var is a valid shell identifier char. var = get_bearer_token_env_var(server.name) assert var == "MCP_LINEAR_APP_TOKEN" @@ -110,7 +109,7 @@ def test_valid_name_derives_shell_safe_env_var(druks_db): async def test_delivery_carries_static_token_in_env(druks_db): - McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) kwargs = await _delivery() assert "linear" in {s.name for s in kwargs["mcp_servers"]} @@ -125,7 +124,7 @@ async def test_required_server_delivers_beside_the_registry(druks_db): # A workspace declares a server with a run-scoped token it minted itself # (Ship's per-repo reviewer token): wire shape + env var ride the same # seam as every registry server. - McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) workspace = _requiring_workspace( RequiredMcpServer( name="github", url="https://api.githubcopilot.com/mcp/", token="ghs_minted" @@ -146,8 +145,8 @@ async def test_required_server_owns_its_name_against_a_registry_twin(druks_db): # is skipped whole: its token neither clobbers the required server's # credential in env nor gets resolved at all (a tokenless twin would # otherwise raise). - McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) - McpServer.create(name="notion", url="https://mcp.notion.com/sse", token="") + await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name="notion", url="https://mcp.notion.com/sse", token="") workspace = _requiring_workspace( RequiredMcpServer( name="linear", url="https://required.internal/linear", token="required-token" @@ -187,14 +186,14 @@ async def test_enabled_static_server_without_token_raises_loudly(druks_db): # A tokenless enabled static row can't authenticate; delivery raises rather # than shipping a header the harness can't fill. (The API rejects creating # one; this guards the model-level path.) - McpServer.create(name="notion", url="https://mcp.notion.com/sse", token="") + await McpServer.create(name="notion", url="https://mcp.notion.com/sse", token="") with pytest.raises(MissingTokenError, match="notion"): await _delivery() async def test_enabled_server_reaches_both_harness_configs_without_token(druks_db): - McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) kwargs = await _delivery() servers = kwargs["mcp_servers"] @@ -224,7 +223,7 @@ async def test_delivery_tolerates_explicit_none_extra_env(druks_db): # ``extra_env=None`` is valid for the underlying run_agent; the fold must treat # it like an omitted env, not unpack None (which would crash the call before it # starts). A static server still rides via its own delivery env. - McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) + await McpServer.create(name="linear", url=_LINEAR_URL, token=_TOKEN) kwargs = await _delivery(extra_env=None) assert "linear" in {s.name for s in kwargs["mcp_servers"]} @@ -234,10 +233,10 @@ async def test_delivery_tolerates_explicit_none_extra_env(druks_db): # --- declared headers: N per server, secret values via env refs ----------- -def _grafana_shaped_server() -> None: +async def _grafana_shaped_server() -> None: # A registry-installed shape: no bearer (empty token_source), one plain # declared header and one secret one. - McpServer.create( + await McpServer.create( name="grafana", url="https://mcp.grafana.com/mcp", token_source="", @@ -247,7 +246,7 @@ def _grafana_shaped_server() -> None: async def test_declared_headers_deliver_inline_and_secret_values_ride_env(druks_db): - _grafana_shaped_server() + await _grafana_shaped_server() kwargs = await _delivery() @@ -265,7 +264,7 @@ async def test_declared_headers_deliver_inline_and_secret_values_ride_env(druks_ async def test_two_header_server_emits_both_headers_in_each_harness_config(druks_db): - _grafana_shaped_server() + await _grafana_shaped_server() kwargs = await _delivery() servers = kwargs["mcp_servers"] header_env_var = servers[0].env_headers["X-Api-Key"] @@ -294,7 +293,7 @@ async def test_two_header_server_emits_both_headers_in_each_harness_config(druks async def test_bearer_and_declared_headers_combine_on_one_server(druks_db): # A static-token server may also declare plain headers; the Authorization # bearer keeps its env-ref form beside them. - McpServer.create( + await McpServer.create( name="acme", url="https://mcp.acme.com/mcp", token=_TOKEN, headers={"X-Region": "eu"} ) @@ -315,7 +314,7 @@ async def test_bearer_and_declared_headers_combine_on_one_server(druks_db): async def test_bearerless_server_delivers_without_a_bearer(druks_db): # The loud MissingTokenError is a static-source contract; a bearerless # server (auth in its headers, or no auth) delivers without any bearer. - McpServer.create(name="public_docs", url="https://docs.example.com/mcp", token_source="") + await McpServer.create(name="public_docs", url="https://docs.example.com/mcp", token_source="") kwargs = await _delivery() @@ -324,10 +323,10 @@ async def test_bearerless_server_delivers_without_a_bearer(druks_db): assert "extra_env" not in kwargs -def test_bearerless_server_merges_with_its_headers(druks_db): - _grafana_shaped_server() +async def test_bearerless_server_merges_with_its_headers(druks_db): + await _grafana_shaped_server() - grafana = McpServer._merged()["grafana"] + grafana = (await McpServer._merged())["grafana"] assert grafana["token_source"] == "" assert grafana["headers"] == {"X-Grafana-URL": "https://acme.grafana.net"} assert grafana["secret_headers"]["X-Api-Key"] == "grafana-api-secret" @@ -336,9 +335,9 @@ def test_bearerless_server_merges_with_its_headers(druks_db): # --- API: CRUD + enable/disable + redaction ------------------------------ -def test_routes_crud_and_token_stays_backend_side(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - created = client.post( +async def test_routes_crud_and_token_stays_backend_side(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + created = await client.post( "/api/mcp-servers", json={"name": "linear", "url": _LINEAR_URL, "token": _TOKEN} ) assert created.status_code == 200 @@ -352,88 +351,89 @@ def test_routes_crud_and_token_stays_backend_side(tmp_path, druks_db): assert "token" not in body # Disable by name, then confirm it drops out of the enabled read path. - toggled = client.patch("/api/mcp-servers/linear", json={"is_enabled": False}) + toggled = await client.patch("/api/mcp-servers/linear", json={"is_enabled": False}) assert toggled.status_code == 200 assert toggled.json()["isEnabled"] is False - listed = client.get("/api/mcp-servers") + listed = await client.get("/api/mcp-servers") assert _TOKEN not in listed.text linear = next(s for s in listed.json() if s["name"] == "linear") assert linear["isEnabled"] is False # Re-adding the same name is rejected — remove first. assert ( - client.post( + await client.post( "/api/mcp-servers", json={"name": "linear", "url": _LINEAR_URL, "token": _TOKEN} - ).status_code - == 409 - ) + ) + ).status_code == 409 - assert client.delete("/api/mcp-servers/linear").status_code == 204 - assert not any(s["name"] == "linear" for s in client.get("/api/mcp-servers").json()) + assert (await client.delete("/api/mcp-servers/linear")).status_code == 204 + assert not any(s["name"] == "linear" for s in (await client.get("/api/mcp-servers")).json()) -def test_routes_reject_invalid_name(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - created = client.post( +async def test_routes_reject_invalid_name(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + created = await client.post( "/api/mcp-servers", json={"name": "linear-app", "url": _LINEAR_URL, "token": _TOKEN} ) assert created.status_code == 422 assert "Invalid MCP server name" in created.text -def test_routes_reject_creating_a_tokenless_custom_server(tmp_path, druks_db): +async def test_routes_reject_creating_a_tokenless_custom_server(tmp_path, druks_db): url = "https://mcp.notion.com/sse" - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: # A custom server is static; a blank (or whitespace-only) token would # create an enabled server that breaks every run at delivery. Rejected at # the boundary instead. for body in ({"name": "notion", "url": url}, {"name": "notion", "url": url, "token": " "}): - created = client.post("/api/mcp-servers", json=body) + created = await client.post("/api/mcp-servers", json=body) assert created.status_code == 422 assert "bearer token" in created.text - assert not any(s["name"] == "notion" for s in client.get("/api/mcp-servers").json()) + assert not any(s["name"] == "notion" for s in (await client.get("/api/mcp-servers")).json()) -def test_routes_reject_creating_a_urlless_custom_server(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: +async def test_routes_reject_creating_a_urlless_custom_server(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: # A blank (or whitespace-only) url is an unreachable endpoint that would # ship into every VM; rejected server-side, not just disabled in the UI. for bad_url in ("", " "): - created = client.post( + created = await client.post( "/api/mcp-servers", json={"name": "notion", "url": bad_url, "token": _TOKEN} ) assert created.status_code == 422 assert "needs a url" in created.text - assert not any(s["name"] == "notion" for s in client.get("/api/mcp-servers").json()) + assert not any(s["name"] == "notion" for s in (await client.get("/api/mcp-servers")).json()) -def test_routes_reject_adding_a_builtin(tmp_path, registry_state, druks_db): +async def test_routes_reject_adding_a_builtin(tmp_path, registry_state, druks_db): load_mcp_catalog(_write_catalog(tmp_path, {"figma_test": _static_entry("https://f/")})) - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: # A catalog entry is built-in — you configure it, you don't add it. - created = client.post( + created = await client.post( "/api/mcp-servers", json={"name": "figma_test", "url": "https://x", "token": "t"} ) assert created.status_code == 409 assert "built-in" in created.text -def test_routes_disable_and_refuse_deleting_a_builtin(tmp_path, registry_state, druks_db): +async def test_routes_disable_and_refuse_deleting_a_builtin(tmp_path, registry_state, druks_db): load_mcp_catalog(_write_catalog(tmp_path, {"figma_test": _static_entry("https://f/")})) - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - figma = next(s for s in client.get("/api/mcp-servers").json() if s["name"] == "figma_test") + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + figma = next( + s for s in (await client.get("/api/mcp-servers")).json() if s["name"] == "figma_test" + ) assert figma["builtin"] is True assert figma["isEnabled"] is True # Backend-owned entry: delete is refused, disable is offered instead. - assert client.delete("/api/mcp-servers/figma_test").status_code == 409 + assert (await client.delete("/api/mcp-servers/figma_test")).status_code == 409 - disabled = client.patch("/api/mcp-servers/figma_test", json={"is_enabled": False}) + disabled = await client.patch("/api/mcp-servers/figma_test", json={"is_enabled": False}) assert disabled.status_code == 200 assert disabled.json()["isEnabled"] is False # The overlay row now exists and the entry reads disabled everywhere. - assert McpServer.get_for_name("figma_test") + assert await McpServer.get_for_name("figma_test") # --- catalog: the deploy-declarative default-server set ------------------- @@ -453,7 +453,7 @@ def _static_entry(url): return {"url": url, "auth": {"type": "static"}} -def test_packaged_catalog_ships_linear_disabled(registry_state, druks_db): +async def test_packaged_catalog_ships_linear_disabled(registry_state, druks_db): # The one packaged default: Linear's hosted MCP, shipped dark — an oauth # entry has no grant until an operator connects it, and an enabled # unconnected one would fail every run's delivery. Ship's github MCP is @@ -461,13 +461,13 @@ def test_packaged_catalog_ships_linear_disabled(registry_state, druks_db): load_mcp_catalog(PACKAGED_MCP_CATALOG) assert "github" not in mcp_servers - builtins = [s for s in McpServer._merged().values() if s["builtin"]] + builtins = [s for s in (await McpServer._merged()).values() if s["builtin"]] assert [s["name"] for s in builtins] == ["linear"] linear = builtins[0] assert linear["url"] == "https://mcp.linear.app/mcp" assert linear["token_source"] == "oauth" assert linear["is_enabled"] is False - assert "linear" not in {s["name"] for s in McpServer.list_enabled()} + assert "linear" not in {s["name"] for s in await McpServer.list_enabled()} async def test_packaged_catalog_delivers_nothing_until_linear_is_connected( @@ -544,19 +544,19 @@ def test_load_catalog_missing_file_fails_loudly(tmp_path): load_mcp_catalog(tmp_path / "absent.json") -def test_db_overlay_still_disables_a_catalog_entry(tmp_path, registry_state, druks_db): +async def test_db_overlay_still_disables_a_catalog_entry(tmp_path, registry_state, druks_db): load_mcp_catalog( _write_catalog(tmp_path, {"figma_test": _static_entry("https://mcp.figma.test/")}) ) - McpServer.create(name="figma_test", url="https://mcp.figma.test/", is_enabled=False) + await McpServer.create(name="figma_test", url="https://mcp.figma.test/", is_enabled=False) - resolved = McpServer._merged()["figma_test"] + resolved = (await McpServer._merged())["figma_test"] assert resolved["builtin"] is True - assert "figma_test" not in {s["name"] for s in McpServer.list_enabled()} + assert "figma_test" not in {s["name"] for s in await McpServer.list_enabled()} -def test_catalog_enabled_false_ships_the_entry_dark(tmp_path, registry_state, druks_db): +async def test_catalog_enabled_false_ships_the_entry_dark(tmp_path, registry_state, druks_db): # ``enabled`` is the catalog's shipped default, not operator state: false # resolves disabled until an operator row says otherwise; an entry without # the key stays enabled exactly as before the field existed. @@ -570,10 +570,10 @@ def test_catalog_enabled_false_ships_the_entry_dark(tmp_path, registry_state, dr ) ) - resolved = McpServer._merged() + resolved = await McpServer._merged() assert resolved["dark_test"]["is_enabled"] is False assert resolved["lit_test"]["is_enabled"] is True - enabled_names = {s["name"] for s in McpServer.list_enabled()} + enabled_names = {s["name"] for s in await McpServer.list_enabled()} assert "dark_test" not in enabled_names assert "lit_test" in enabled_names @@ -614,7 +614,7 @@ async def test_definition_auth_wins_over_an_overlay_row_token( # decides how the token is sourced — a row token is inert for env-sourced # entries, and druks never needs one stored. load_mcp_catalog(_write_catalog(tmp_path, {"vault_test": _env_entry()})) - McpServer.create(name="vault_test", url="https://mcp.vault.test/", token="db-token") + await McpServer.create(name="vault_test", url="https://mcp.vault.test/", token="db-token") monkeypatch.setenv("VAULT_TEST_TOKEN", "env-token") kwargs = await _delivery() @@ -622,20 +622,22 @@ async def test_definition_auth_wins_over_an_overlay_row_token( assert kwargs["extra_env"][get_bearer_token_env_var("vault_test")] == "env-token" -def test_api_has_token_reflects_env_presence_for_env_sourced( +async def test_api_has_token_reflects_env_presence_for_env_sourced( tmp_path, registry_state, monkeypatch, druks_db ): load_mcp_catalog(_write_catalog(tmp_path, {"vault_test": _env_entry()})) monkeypatch.delenv("VAULT_TEST_TOKEN", raising=False) - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - vault = next(s for s in client.get("/api/mcp-servers").json() if s["name"] == "vault_test") + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + vault = next( + s for s in (await client.get("/api/mcp-servers")).json() if s["name"] == "vault_test" + ) assert vault["hasToken"] is False # The badge can name the var to set — a var name, never a value. assert vault["sourceEnvVar"] == "VAULT_TEST_TOKEN" monkeypatch.setenv("VAULT_TEST_TOKEN", "vault-secret") - listed = client.get("/api/mcp-servers") + listed = await client.get("/api/mcp-servers") vault = next(s for s in listed.json() if s["name"] == "vault_test") assert vault["hasToken"] is True assert "vault-secret" not in listed.text diff --git a/backend/tests/test_notification_destinations.py b/backend/tests/test_notification_destinations.py index 62e12712..c6fdcb33 100644 --- a/backend/tests/test_notification_destinations.py +++ b/backend/tests/test_notification_destinations.py @@ -11,8 +11,7 @@ UnknownDestinationKindError, ) from druks.notifications.models import Destination, DestinationKind -from druks.testing import configure_app_for_test, make_settings -from fastapi.testclient import TestClient +from druks.testing import asgi_client, configure_app_for_test, make_settings _WEBHOOK_URL = "https://hooks.slack.com/services/T000/B000/secretpart" @@ -91,29 +90,29 @@ def fake_slack(monkeypatch): # --- registry: CRUD round-trip + kind gate -------------------------------- -def test_create_get_list_delete_round_trip(druks_db): - beta = Destination.create(name="beta", kind="slack_webhook", url=_WEBHOOK_URL) - alpha = Destination.create(name="alpha", kind="slack_webhook", url=_WEBHOOK_URL) +async def test_create_get_list_delete_round_trip(druks_db): + beta = await Destination.create(name="beta", kind="slack_webhook", url=_WEBHOOK_URL) + alpha = await Destination.create(name="alpha", kind="slack_webhook", url=_WEBHOOK_URL) - assert Destination.get(beta.id).id == beta.id - assert Destination.get_for_name("alpha").id == alpha.id - assert Destination.get("no-such-id") is None - assert Destination.get_for_name("no-such-name") is None - assert [destination.name for destination in Destination.list_all()] == ["alpha", "beta"] + assert (await Destination.get(beta.id)).id == beta.id + assert (await Destination.get_for_name("alpha")).id == alpha.id + assert await Destination.get("no-such-id") is None + assert await Destination.get_for_name("no-such-name") is None + assert [destination.name for destination in await Destination.list_all()] == ["alpha", "beta"] assert beta.is_enabled is True - beta.delete() - assert Destination.get_for_name("beta") is None - assert [destination.name for destination in Destination.list_all()] == ["alpha"] + await beta.delete() + assert await Destination.get_for_name("beta") is None + assert [destination.name for destination in await Destination.list_all()] == ["alpha"] -def test_create_rejects_unknown_kind_without_echoing_the_url(druks_db): +async def test_create_rejects_unknown_kind_without_echoing_the_url(druks_db): with pytest.raises(UnknownDestinationKindError) as excinfo: - Destination.create(name="pager", kind="pagerduty", url=_WEBHOOK_URL) + await Destination.create(name="pager", kind="pagerduty", url=_WEBHOOK_URL) assert "pagerduty" in str(excinfo.value) assert _WEBHOOK_URL not in str(excinfo.value) - assert Destination.get_for_name("pager") is None + assert await Destination.get_for_name("pager") is None # --- routes: CRUD + redaction --------------------------------------------- @@ -125,9 +124,9 @@ def _create_body(**overrides) -> dict: return body -def test_routes_create_masks_url(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - created = client.post("/api/notifications/destinations", json=_create_body()) +async def test_routes_create_masks_url(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + created = await client.post("/api/notifications/destinations", json=_create_body()) assert created.status_code == 200 body = created.json() @@ -140,81 +139,89 @@ def test_routes_create_masks_url(tmp_path, druks_db): assert "secretpart" not in created.text -def test_routes_reject_duplicate_name(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - first = client.post("/api/notifications/destinations", json=_create_body()) +async def test_routes_reject_duplicate_name(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + first = await client.post("/api/notifications/destinations", json=_create_body()) assert first.status_code == 200 - duplicate = client.post("/api/notifications/destinations", json=_create_body()) + duplicate = await client.post("/api/notifications/destinations", json=_create_body()) assert duplicate.status_code == 409 assert "secretpart" not in duplicate.text -def test_routes_reject_unknown_kind(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - created = client.post( +async def test_routes_reject_unknown_kind(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + created = await client.post( "/api/notifications/destinations", json=_create_body(kind="pagerduty") ) assert created.status_code == 422 assert "secretpart" not in created.text - assert not client.get("/api/notifications/destinations").json() + assert not (await client.get("/api/notifications/destinations")).json() -def test_routes_reject_undeliverable_url(tmp_path, druks_db): +async def test_routes_reject_undeliverable_url(tmp_path, druks_db): # Save-time deliverability: the same offline apprise parse the send path # uses, so a typo fails while the operator is present — not at first park. - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: for bad_url in ("https://example.com/hook", "not-a-url"): - created = client.post("/api/notifications/destinations", json=_create_body(url=bad_url)) + created = await client.post( + "/api/notifications/destinations", json=_create_body(url=bad_url) + ) assert created.status_code == 422 assert bad_url not in created.text - assert not client.get("/api/notifications/destinations").json() + assert not (await client.get("/api/notifications/destinations")).json() -def test_routes_reject_blank_name(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: +async def test_routes_reject_blank_name(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: for blank in ("", " "): - created = client.post("/api/notifications/destinations", json=_create_body(name=blank)) + created = await client.post( + "/api/notifications/destinations", json=_create_body(name=blank) + ) assert created.status_code == 422 assert "secretpart" not in created.text - assert not client.get("/api/notifications/destinations").json() + assert not (await client.get("/api/notifications/destinations")).json() -def test_routes_toggle_enabled(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - destination_id = client.post("/api/notifications/destinations", json=_create_body()).json()[ - "id" - ] +async def test_routes_toggle_enabled(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + destination_id = ( + await client.post("/api/notifications/destinations", json=_create_body()) + ).json()["id"] - toggled = client.patch( + toggled = await client.patch( f"/api/notifications/destinations/{destination_id}", json={"is_enabled": False} ) assert toggled.status_code == 200 assert toggled.json()["isEnabled"] is False - listed = client.get("/api/notifications/destinations").json() + listed = (await client.get("/api/notifications/destinations")).json() assert listed[0]["isEnabled"] is False - back_on = client.patch( + back_on = await client.patch( f"/api/notifications/destinations/{destination_id}", json={"is_enabled": True} ) assert back_on.json()["isEnabled"] is True - missing = client.patch( + missing = await client.patch( "/api/notifications/destinations/no-such-id", json={"is_enabled": False} ) assert missing.status_code == 404 -def test_routes_delete(tmp_path, druks_db): - with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - destination_id = client.post("/api/notifications/destinations", json=_create_body()).json()[ - "id" - ] +async def test_routes_delete(tmp_path, druks_db): + async with asgi_client(configure_app_for_test(settings=make_settings(tmp_path))) as client: + destination_id = ( + await client.post("/api/notifications/destinations", json=_create_body()) + ).json()["id"] - assert client.delete(f"/api/notifications/destinations/{destination_id}").status_code == 204 - assert not client.get("/api/notifications/destinations").json() - assert client.delete(f"/api/notifications/destinations/{destination_id}").status_code == 404 + assert ( + await client.delete(f"/api/notifications/destinations/{destination_id}") + ).status_code == 204 + assert not (await client.get("/api/notifications/destinations")).json() + assert ( + await client.delete(f"/api/notifications/destinations/{destination_id}") + ).status_code == 404 # --- informational delivery (Apprise) -------------------------------------- diff --git a/backend/tests/test_notifications.py b/backend/tests/test_notifications.py index 2895dfd9..1b0392a1 100644 --- a/backend/tests/test_notifications.py +++ b/backend/tests/test_notifications.py @@ -15,16 +15,16 @@ _SUBJECT = {"type": "work_item", "id": 1} -def _destination(name: str = "ops") -> Destination: - return Destination.create(name=name, kind="slack_webhook", url=_WEBHOOK_URL) +async def _destination(name: str = "ops") -> Destination: + return await Destination.create(name=name, kind="slack_webhook", url=_WEBHOOK_URL) # --- entity ------------------------------------------------------------------ -def test_unique_token_collision_raises(druks_db): - destination = _destination() - first = Notification.create( +async def test_unique_token_collision_raises(druks_db): + destination = await _destination() + first = await Notification.create( destination_id=destination.id, reason="r", body="b", subject=_SUBJECT ) @@ -37,40 +37,42 @@ def test_unique_token_collision_raises(druks_db): ) druks_db.add(duplicate) with pytest.raises(IntegrityError): - druks_db.flush() + await druks_db.flush() -def test_list_recent_newest_first_with_limit(druks_db): - destination = _destination() +async def test_list_recent_newest_first_with_limit(druks_db): + destination = await _destination() ids = [ - Notification.create( - destination_id=destination.id, reason="r", body=f"note {i}", subject=_SUBJECT + ( + await Notification.create( + destination_id=destination.id, reason="r", body=f"note {i}", subject=_SUBJECT + ) ).id for i in range(3) ] - listed = Notification.list_recent(limit=2) + listed = await Notification.list_recent(limit=2) assert [notification.id for notification in listed] == [ids[2], ids[1]] # --- read endpoints --------------------------------------------------------- -def test_endpoints_list_and_get_omit_the_token(tmp_path, druks_db): - destination = _destination() +async def test_endpoints_list_and_get_omit_the_token(tmp_path, druks_db): + destination = await _destination() tokens = [] for index in range(3): - notification = Notification.create( + notification = await Notification.create( destination_id=destination.id, reason="gate.parked", body=f"note {index}", subject=_SUBJECT, ) tokens.append(notification.correlation_token) - failed = Notification.create( + failed = await Notification.create( destination_id=destination.id, reason="gate.parked", body="bad", subject=_SUBJECT ) - failed.mark_failed("DeliveryError: HTTPStatusError") + await failed.mark_failed("DeliveryError: HTTPStatusError") tokens.append(failed.correlation_token) with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: @@ -97,10 +99,10 @@ def test_endpoints_list_and_get_omit_the_token(tmp_path, druks_db): assert client.get("/api/notifications/no-such-id").status_code == 404 -def test_destinations_route_still_resolves_after_notifications_mount(tmp_path, druks_db): +async def test_destinations_route_still_resolves_after_notifications_mount(tmp_path, druks_db): # The route-order pin: the notifications /{notification_id} match must not # swallow /api/notifications/destinations. - _destination(name="alpha") + await _destination(name="alpha") with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: listed = client.get("/api/notifications/destinations") @@ -112,8 +114,8 @@ def test_destinations_route_still_resolves_after_notifications_mount(tmp_path, d # --- the gate-park destination setting --------------------------------------- -def test_settings_gate_park_destination_set_clear_and_reject(tmp_path, druks_db): - destination = _destination() +async def test_settings_gate_park_destination_set_clear_and_reject(tmp_path, druks_db): + destination = await _destination() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: assert client.get("/api/settings").json()["gateParkDestinationId"] is None @@ -131,8 +133,8 @@ def test_settings_gate_park_destination_set_clear_and_reject(tmp_path, druks_db) assert cleared.json()["gateParkDestinationId"] is None -def test_deleting_designated_destination_unsets_the_pointer(tmp_path, druks_db): - destination = _destination() +async def test_deleting_designated_destination_unsets_the_pointer(tmp_path, druks_db): + destination = await _destination() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: client.patch("/api/settings", json={"gateParkDestinationId": destination.id}) @@ -151,7 +153,7 @@ def test_deleting_designated_destination_unsets_the_pointer(tmp_path, druks_db): } -def _parked_notification(druks_db, *, ask=None, run_state="parked"): +async def _parked_notification(druks_db, *, ask=None, run_state="parked"): ask = ask or _IN_APP_ASK run = Run( id=str(uuid7()), @@ -161,10 +163,10 @@ def _parked_notification(druks_db, *, ask=None, run_state="parked"): input_requested_at=Base.utc_now(), ) druks_db.add(run) - druks_db.flush() - seed_dbos_status(druks_db, run.id, run_state) - destination = _destination(name=f"dest-{run.id[-8:]}") - notification = Notification.create( + await druks_db.flush() + await seed_dbos_status(druks_db, run.id, run_state) + destination = await _destination(name=f"dest-{run.id[-8:]}") + notification = await Notification.create( destination_id=destination.id, reason="gate.parked", body="review the plan", @@ -187,7 +189,7 @@ async def _spy(self, **fields): async def test_respond_resumes_and_marks_acknowledged(druks_db, resume_spy): - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) await respond_to_notification( notification.correlation_token, @@ -195,12 +197,12 @@ async def test_respond_resumes_and_marks_acknowledged(druks_db, resume_spy): ) assert resume_spy == [{"id": run.id, "action": "approve", "answers": {"q1": "a"}, "note": ""}] - assert Notification.get(notification.id).state == "acknowledged" - assert Notification.get(notification.id).is_acknowledged + assert (await Notification.get(notification.id)).state == "acknowledged" + assert (await Notification.get(notification.id)).is_acknowledged async def test_respond_route_codes_and_secret_hygiene(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) token = notification.correlation_token client = TestClient(configure_app_for_test(settings=make_settings(tmp_path))) @@ -246,7 +248,7 @@ async def test_respond_route_codes_and_secret_hygiene(tmp_path, druks_db, resume async def test_respond_external_notification_not_answerable(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification( + run, notification = await _parked_notification( druks_db, ask={"presentation": "external", "label": "Answer on the ticket"} ) client = TestClient(configure_app_for_test(settings=make_settings(tmp_path))) @@ -262,8 +264,8 @@ async def test_respond_external_notification_not_answerable(tmp_path, druks_db, async def test_respond_runless_notification_not_answerable(tmp_path, druks_db, resume_spy): - destination = _destination(name="runless-dest") - notification = Notification.create( + destination = await _destination(name="runless-dest") + notification = await Notification.create( destination_id=destination.id, reason="r", body="b", subject=_SUBJECT ) client = TestClient(configure_app_for_test(settings=make_settings(tmp_path))) @@ -286,10 +288,10 @@ async def test_respond_stale_round_409(tmp_path, druks_db, resume_spy): input_requested_at=Base.utc_now(), ) druks_db.add(run) - druks_db.flush() - seed_dbos_status(druks_db, run.id, "parked") - destination = _destination(name="stale-dest") - notification = Notification.create( + await druks_db.flush() + await seed_dbos_status(druks_db, run.id, "parked") + destination = await _destination(name="stale-dest") + notification = await Notification.create( destination_id=destination.id, reason="gate.parked", body="review", @@ -309,7 +311,7 @@ async def test_respond_stale_round_409(tmp_path, druks_db, resume_spy): async def test_respond_run_no_longer_parked_409(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db, run_state="finished") + run, notification = await _parked_notification(druks_db, run_state="finished") client = TestClient(configure_app_for_test(settings=make_settings(tmp_path))) response = client.post( @@ -324,8 +326,12 @@ async def test_respond_run_no_longer_parked_409(tmp_path, druks_db, resume_spy): async def test_respond_corrupt_correlation_500_and_logged( tmp_path, druks_db, resume_spy, monkeypatch, caplog ): - run, notification = _parked_notification(druks_db) - monkeypatch.setattr(Run, "get", classmethod(lambda cls, run_id: None)) + run, notification = await _parked_notification(druks_db) + + async def _missing(cls, run_id): + return None + + monkeypatch.setattr(Run, "get", classmethod(_missing)) client = TestClient( configure_app_for_test(settings=make_settings(tmp_path)), raise_server_exceptions=False ) @@ -345,7 +351,7 @@ async def test_respond_corrupt_correlation_500_and_logged( async def test_respond_direct_call_rejects_whitespace_only_content(druks_db, resume_spy): # The core is also the direct-call boundary (the Slack rail bypasses the # HTTP models' whitespace stripping) — blank means blank on every path. - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) with pytest.raises(InvalidChoiceError): await respond_to_notification( @@ -358,13 +364,13 @@ async def test_respond_direct_call_rejects_whitespace_only_content(druks_db, res {"control": "request_changes", "note": " "}, ) assert resume_spy == [] - assert Notification.get(notification.id).state == "pending" + assert (await Notification.get(notification.id)).state == "pending" async def test_respond_ask_without_presentation_not_answerable(tmp_path, druks_db, resume_spy): # An ask that doesn't declare in_app isn't answerable via this rail — the # mapped 422, not a crash. - run, notification = _parked_notification( + run, notification = await _parked_notification( druks_db, ask={"label": "legacy ask", "controls": ["approve"]} ) client = TestClient(configure_app_for_test(settings=make_settings(tmp_path))) diff --git a/backend/tests/test_notifications_durable.py b/backend/tests/test_notifications_durable.py index 84f8c99b..97a76e02 100644 --- a/backend/tests/test_notifications_durable.py +++ b/backend/tests/test_notifications_durable.py @@ -20,7 +20,9 @@ from druks.workflows import Gate, OperatorReply, Run, Workflow from fastapi.testclient import TestClient from pydantic import BaseModel, Field -from sqlalchemy import create_engine, select, text +from sqlalchemy import NullPool, create_engine, select, text +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy.orm import Session PG_BASE = os.environ.get("DRUKS_TEST_PG", "postgresql://druks:druks@localhost:5432") DB = "druks_notifications_durable_test" @@ -119,7 +121,7 @@ async def run_multistep(self) -> None: @pytest.fixture(scope="module", autouse=True) -def rt(): +async def rt(): db_url_snap = os.environ.get("DRUKS_DATABASE_URL") admin = psycopg.connect(f"{PG_BASE}/postgres", autocommit=True) @@ -127,19 +129,19 @@ def rt(): admin.execute(f"CREATE DATABASE {DB}") admin.close() - engine = create_engine(URL) - init_db(engine) - configure_engine(engine) - configure_session(engine) - session = get_session(engine) - try: + schema_engine = create_engine(URL) + init_db(schema_engine) + with Session(schema_engine) as session: session.add_all( NotificationProbe(id=subject_id) for subject_id in (9001, 9002, 9003, 9004, 9005, 9006, 9007, 9008, 9009, 9010, 9014) ) session.commit() - finally: - session.close() + # NullPool: the TestClient below runs requests on its own loop, and a pooled + # async connection must never cross loops. + engine = create_async_engine(URL, poolclass=NullPool) + configure_engine(engine) + configure_session(engine) ( in_app_flow, external_flow, @@ -151,10 +153,11 @@ def rt(): # The outbox module was imported above — its queue + workflow register # before launch(), which is the wiring this whole module runs through. init_dbos() - launch() + await launch() try: yield SimpleNamespace( engine=engine, + schema_engine=schema_engine, InAppFlow=in_app_flow, ExternalFlow=external_flow, SubjectlessFlow=subjectless_flow, @@ -163,7 +166,8 @@ def rt(): ) finally: shutdown() - engine.dispose() + await engine.dispose() + schema_engine.dispose() for kind in ( "in_app_flow", "external_flow", @@ -203,52 +207,59 @@ async def __call__(self, destination, body, *, actions=None, token=None, idempot @pytest.fixture -def deliver_spy(monkeypatch): +async def deliver_spy(monkeypatch): spy = _DeliverSpy() monkeypatch.setattr(outbox, "deliver", spy) return spy -def _seed(rt, seeder): +async def _seed(rt, seeder): # Commit for real: the outbox worker reads through its own sessions. # Expunge first so the returned instance keeps its loaded attributes past # the commit's expiry. session = get_session(rt.engine) db_session.registry.set(session) try: - result = seeder() - session.flush() + result = await seeder() + await session.flush() session.expunge_all() - session.commit() + await session.commit() return result finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() + + +async def _seed_destination(rt, name): + async def create(): + return await Destination.create(name=name, kind="slack_webhook", url=_WEBHOOK_URL) + + return await _seed(rt, create) async def _deliver(rt, *, to, subject=None, reason="r", body="b", actions=None): # The create-seam path a producer uses (the gate-park producer's shape): # persist the row committed, then enqueue the outbox — no notify() hatch. - notification_id = _seed( - rt, - lambda: ( - Notification.create( - destination_id=Destination.get_for_name(to).id, - reason=reason, - body=body, - subject=subject or {"type": "notification_probe", "id": 1}, - actions=actions, - ).id - ), - ) + async def create(): + destination = await Destination.get_for_name(to) + notification = await Notification.create( + destination_id=destination.id, + reason=reason, + body=body, + subject=subject or {"type": "notification_probe", "id": 1}, + actions=actions, + ) + return notification.id + + notification_id = await _seed(rt, create) await notifications_queue.enqueue_async(send_notification, notification_id) return notification_id -def _snapshot(rt, notification_id) -> dict: +async def _snapshot(rt, notification_id) -> dict: session = get_session(rt.engine) try: - notification = session.get(Notification, notification_id) + notification = await session.get(Notification, notification_id) return { "state": notification.state, "attempts": notification.attempts, @@ -257,21 +268,21 @@ def _snapshot(rt, notification_id) -> dict: "token": notification.correlation_token, } finally: - session.close() + await session.close() async def _wait_for(rt, notification_id, predicate, timeout=30.0): deadline = asyncio.get_event_loop().time() + timeout while asyncio.get_event_loop().time() < deadline: - snapshot = _snapshot(rt, notification_id) + snapshot = await _snapshot(rt, notification_id) if predicate(snapshot): return snapshot await asyncio.sleep(0.1) - raise AssertionError(f"timed out; last={_snapshot(rt, notification_id)}") + raise AssertionError(f"timed out; last={await _snapshot(rt, notification_id)}") async def test_outbox_delivers_with_actions_token_and_key(rt, deliver_spy): - _seed(rt, lambda: Destination.create(name="happy", kind="slack_webhook", url=_WEBHOOK_URL)) + await _seed_destination(rt, "happy") notification_id = await _deliver( rt, to="happy", reason="ops.alert", body="hello", actions=[{"id": "ok", "label": "OK"}] @@ -290,7 +301,7 @@ async def test_outbox_delivers_with_actions_token_and_key(rt, deliver_spy): async def test_transient_failure_retries_to_delivered_and_touches_no_run(rt, deliver_spy): deliver_spy.failures_remaining = 2 - _seed(rt, lambda: Destination.create(name="flaky", kind="slack_webhook", url=_WEBHOOK_URL)) + await _seed_destination(rt, "flaky") notification_id = await _deliver(rt, to="flaky") @@ -298,7 +309,7 @@ async def test_transient_failure_retries_to_delivered_and_touches_no_run(rt, del assert done["attempts"] == 3 assert len(deliver_spy.calls) == 3 # Delivery is decoupled from the run lifecycle: no Run row exists or was touched. - with rt.engine.connect() as connection: + with rt.schema_engine.connect() as connection: assert connection.execute(text("SELECT count(*) FROM durable_runs")).scalar_one() == 0 @@ -309,7 +320,7 @@ async def test_terminal_failure_marks_failed_sanitized_and_reads_back( deliver_spy.error = DeliveryError("doomed", "HTTPStatusError") # Two attempts instead of five: same terminal path, fraction of the backoff. monkeypatch.setattr(outbox, "_SEND_RETRIES", {"retries_allowed": True, "max_attempts": 2}) - _seed(rt, lambda: Destination.create(name="doomed", kind="slack_webhook", url=_WEBHOOK_URL)) + await _seed_destination(rt, "doomed") notification_id = await _deliver(rt, to="doomed") @@ -334,7 +345,7 @@ async def test_unexpected_error_reduces_to_class_name(rt, deliver_spy, monkeypat deliver_spy.always_fail = True deliver_spy.error = RuntimeError(f"boom at {_WEBHOOK_URL}") monkeypatch.setattr(outbox, "_SEND_RETRIES", {"retries_allowed": True, "max_attempts": 2}) - _seed(rt, lambda: Destination.create(name="leaky", kind="slack_webhook", url=_WEBHOOK_URL)) + await _seed_destination(rt, "leaky") notification_id = await _deliver(rt, to="leaky") @@ -344,7 +355,7 @@ async def test_unexpected_error_reduces_to_class_name(rt, deliver_spy, monkeypat async def test_rerun_on_delivered_notification_skips_the_send(rt, deliver_spy): - _seed(rt, lambda: Destination.create(name="once", kind="slack_webhook", url=_WEBHOOK_URL)) + await _seed_destination(rt, "once") notification_id = await _deliver(rt, to="once") await _wait_for(rt, notification_id, lambda s: s["state"] == "delivered") assert len(deliver_spy.calls) == 1 @@ -356,96 +367,91 @@ async def test_rerun_on_delivered_notification_skips_the_send(rt, deliver_spy): async def test_create_seam_plus_direct_enqueue_delivers(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="seam", kind="slack_webhook", url=_WEBHOOK_URL) - ) - notification = _seed( - rt, - lambda: Notification.create( - destination_id=destination.id, reason="r", body="b", subject={"type": "probe", "id": 1} - ), - ) + destination = await _seed_destination(rt, "seam") + + async def create(): + return await Notification.create( + destination_id=destination.id, + reason="r", + body="b", + subject={"type": "probe", "id": 1}, + ) + + notification = await _seed(rt, create) # The create seam persisted a pending row and enqueued nothing. - assert _snapshot(rt, notification.id)["state"] == "pending" + assert (await _snapshot(rt, notification.id))["state"] == "pending" assert deliver_spy.calls == [] handle = await notifications_queue.enqueue_async(send_notification, notification.id) await handle.get_result() - assert _snapshot(rt, notification.id)["state"] == "delivered" + assert (await _snapshot(rt, notification.id))["state"] == "delivered" assert len(deliver_spy.calls) == 1 # --- gate-park notifications --------------------------------------------------- -def _set_gate_park_pointer(rt, destination_id): +async def _set_gate_park_pointer(rt, destination_id): session = get_session(rt.engine) db_session.registry.set(session) try: - UserSettings.get().set_gate_park_destination(destination_id) - session.commit() + await (await UserSettings.get()).set_gate_park_destination(destination_id) + await session.commit() finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() -def _run_snapshot(rt, workflow_id) -> Run: +async def _run_snapshot(rt, workflow_id) -> Run: session = get_session(rt.engine) try: - return session.get(Run, workflow_id) + return await session.get(Run, workflow_id) finally: - session.close() + await session.close() -def _notifications_for_run(rt, workflow_id) -> list[Notification]: +async def _notifications_for_run(rt, workflow_id) -> list[Notification]: session = get_session(rt.engine) try: - return list( - session.execute( - select(Notification) - .where(Notification.run_id == workflow_id) - .order_by(Notification.id) - ) - .scalars() - .all() + result = await session.execute( + select(Notification).where(Notification.run_id == workflow_id).order_by(Notification.id) ) + return list(result.scalars().all()) finally: - session.close() + await session.close() async def _wait_run(rt, workflow_id, predicate, timeout=30.0) -> Run: deadline = asyncio.get_event_loop().time() + timeout while asyncio.get_event_loop().time() < deadline: - run = _run_snapshot(rt, workflow_id) + run = await _run_snapshot(rt, workflow_id) if run and predicate(run): return run await asyncio.sleep(0.1) - raise AssertionError(f"timed out; last={_run_snapshot(rt, workflow_id)}") + raise AssertionError(f"timed out; last={await _run_snapshot(rt, workflow_id)}") async def _wait_notification(rt, workflow_id, state, timeout=30.0) -> Notification: deadline = asyncio.get_event_loop().time() + timeout while asyncio.get_event_loop().time() < deadline: - rows = _notifications_for_run(rt, workflow_id) + rows = await _notifications_for_run(rt, workflow_id) if rows and rows[0].state == state: return rows[0] await asyncio.sleep(0.1) - raise AssertionError(f"timed out; last={_notifications_for_run(rt, workflow_id)}") + raise AssertionError(f"timed out; last={await _notifications_for_run(rt, workflow_id)}") async def test_in_app_park_notifies_with_actions(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-inapp", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-inapp") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.InAppFlow.start(subject=NotificationProbe(id=9001)) parked = await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) notification = await _wait_notification(rt, workflow_id, "delivered") - assert len(_notifications_for_run(rt, workflow_id)) == 1 + assert len(await _notifications_for_run(rt, workflow_id)) == 1 assert notification.reason == "gate.parked" assert notification.body.startswith("Review") assert "Which database?" in notification.body @@ -465,10 +471,8 @@ async def test_in_app_park_notifies_with_actions(rt, deliver_spy): async def test_external_park_notifies_without_actions(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-ext", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-ext") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.ExternalFlow.start(subject=NotificationProbe(id=9002)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) @@ -482,10 +486,8 @@ async def test_external_park_notifies_without_actions(rt, deliver_spy): async def test_external_park_with_declared_url_sets_deep_link(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-url", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-url") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.ExternalUrlFlow.start(subject=NotificationProbe(id=9003)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) @@ -496,7 +498,7 @@ async def test_external_park_with_declared_url_sets_deep_link(rt, deliver_spy): async def test_no_designated_destination_notifies_nothing(rt, deliver_spy): - _set_gate_park_pointer(rt, None) + await _set_gate_park_pointer(rt, None) in_app_id = await rt.InAppFlow.start(subject=NotificationProbe(id=9004)) external_id = await rt.ExternalFlow.start(subject=NotificationProbe(id=9014)) @@ -504,54 +506,52 @@ async def test_no_designated_destination_notifies_nothing(rt, deliver_spy): await _wait_run(rt, external_id, lambda run: run.state == RunState.PARKED) await asyncio.sleep(1.0) - assert _notifications_for_run(rt, in_app_id) == [] - assert _notifications_for_run(rt, external_id) == [] + assert await _notifications_for_run(rt, in_app_id) == [] + assert await _notifications_for_run(rt, external_id) == [] assert deliver_spy.calls == [] async def test_deleted_designated_destination_notifies_nothing(rt, deliver_spy): - destination = _seed( - rt, - lambda: Destination.create(name="inbox-deleted", kind="slack_webhook", url=_WEBHOOK_URL), - ) - _set_gate_park_pointer(rt, destination.id) - _seed(rt, lambda: Destination.get(destination.id).delete()) + destination = await _seed_destination(rt, "inbox-deleted") + await _set_gate_park_pointer(rt, destination.id) + + async def delete_destination(): + await (await Destination.get(destination.id)).delete() + + await _seed(rt, delete_destination) workflow_id = await rt.ExternalFlow.start(subject=NotificationProbe(id=9005)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) await asyncio.sleep(1.0) - assert _notifications_for_run(rt, workflow_id) == [] + assert await _notifications_for_run(rt, workflow_id) == [] assert deliver_spy.calls == [] # ON DELETE SET NULL cleared the pointer itself. session = get_session(rt.engine) try: - assert session.get(UserSettings, UserSettings.SINGLETON_ID).gate_park_destination_id is None + settings = await session.get(UserSettings, UserSettings.SINGLETON_ID) + assert settings.gate_park_destination_id is None finally: - session.close() + await session.close() async def test_subjectless_park_notifies_nothing(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-nosubj", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-nosubj") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.SubjectlessFlow.start(subject=None) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) await asyncio.sleep(1.0) - assert _notifications_for_run(rt, workflow_id) == [] + assert await _notifications_for_run(rt, workflow_id) == [] assert deliver_spy.calls == [] async def test_failed_delivery_leaves_run_parked_and_resumable(rt, deliver_spy, monkeypatch): deliver_spy.always_fail = True monkeypatch.setattr(outbox, "_SEND_RETRIES", {"retries_allowed": True, "max_attempts": 2}) - destination = _seed( - rt, lambda: Destination.create(name="inbox-flaky", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-flaky") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.InAppFlow.start(subject=NotificationProbe(id=9006)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) @@ -559,17 +559,15 @@ async def test_failed_delivery_leaves_run_parked_and_resumable(rt, deliver_spy, assert _WEBHOOK_URL not in notification.last_error # The park never noticed the dead endpoint: still waiting, still resumable. - parked = _run_snapshot(rt, workflow_id) + parked = await _run_snapshot(rt, workflow_id) assert parked.state == RunState.PARKED await parked.resume(action="approve") await _wait_run(rt, workflow_id, lambda run: run.state == RunState.FINISHED) async def test_replayed_park_notifies_once(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-replay", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-replay") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.ExternalFlow.start(subject=NotificationProbe(id=9007)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) @@ -577,7 +575,7 @@ async def test_replayed_park_notifies_once(rt, deliver_spy): assert len(deliver_spy.calls) == 1 def _execution_claim(): - with rt.engine.connect() as connection: + with rt.schema_engine.connect() as connection: return connection.execute( text( "SELECT started_at_epoch_ms FROM dbos.workflow_status WHERE workflow_uuid = :id" @@ -602,7 +600,7 @@ def _execution_claim(): await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) await asyncio.sleep(1.0) # room for a wrong duplicate enqueue to land - assert len(_notifications_for_run(rt, workflow_id)) == 1 + assert len(await _notifications_for_run(rt, workflow_id)) == 1 assert len(deliver_spy.calls) == 1 # Two recv waiters now share the topic (the pre-resume one and the @@ -610,18 +608,17 @@ def _execution_claim(): session = get_session(rt.engine) db_session.registry.set(session) try: - await session.get(Run, workflow_id).cancel() - session.commit() + run = await session.get(Run, workflow_id) + await run.cancel() + await session.commit() finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() async def test_each_park_round_gets_its_own_notification(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-rounds", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-rounds") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.DoubleParkFlow.start(subject=NotificationProbe(id=9008)) first_round = await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) @@ -641,14 +638,14 @@ async def test_each_park_round_gets_its_own_notification(rt, deliver_spy): ) deadline = asyncio.get_event_loop().time() + 30 while asyncio.get_event_loop().time() < deadline: - rows = _notifications_for_run(rt, workflow_id) + rows = await _notifications_for_run(rt, workflow_id) if len(rows) == 2 and rows[1].state == "delivered": break await asyncio.sleep(0.1) else: raise AssertionError(f"second round never notified; {deliver_spy.calls=}") - rows = _notifications_for_run(rt, workflow_id) + rows = await _notifications_for_run(rt, workflow_id) assert [row.body for row in rows] == ["Round one", "Round two"] assert rows[0].run_parked_at != rows[1].run_parked_at assert len(deliver_spy.calls) == 2 @@ -666,18 +663,18 @@ async def _respond_in_own_session(rt, token, choice): db_session.registry.set(session) try: await respond_to_notification(token, choice) - session.commit() + await session.commit() return "ok" except NotificationError as error: - session.rollback() + await session.rollback() return type(error).__name__ finally: - db_session.remove() - session.close() + await db_session.remove() + await session.close() def _dbos_replies(rt, workflow_id) -> int: - with rt.engine.connect() as connection: + with rt.schema_engine.connect() as connection: return connection.execute( text( "SELECT count(*) FROM dbos.notifications" @@ -688,10 +685,8 @@ def _dbos_replies(rt, workflow_id) -> int: async def test_respond_round_trip_finishes_the_run(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-respond", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-respond") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.InAppFlow.start(subject=NotificationProbe(id=9009)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) notification = await _wait_notification(rt, workflow_id, "delivered") @@ -704,7 +699,7 @@ async def test_respond_round_trip_finishes_the_run(rt, deliver_spy): await _wait_run(rt, workflow_id, lambda run: run.state == RunState.FINISHED) assert OperatorReply(action="approve", answers={"q1": "postgres"}) in _REVIEW_REPLIES - assert _notifications_for_run(rt, workflow_id)[0].state == "acknowledged" + assert (await _notifications_for_run(rt, workflow_id))[0].state == "acknowledged" # Sequential second answer: the acknowledged fast-path rejects it, and the # round's DBOS bookkeeping still holds exactly one reply. @@ -714,10 +709,8 @@ async def test_respond_round_trip_finishes_the_run(rt, deliver_spy): async def test_concurrent_responds_resolve_to_one_answer(rt, deliver_spy): - destination = _seed( - rt, lambda: Destination.create(name="inbox-race", kind="slack_webhook", url=_WEBHOOK_URL) - ) - _set_gate_park_pointer(rt, destination.id) + destination = await _seed_destination(rt, "inbox-race") + await _set_gate_park_pointer(rt, destination.id) workflow_id = await rt.InAppFlow.start(subject=NotificationProbe(id=9010)) await _wait_run(rt, workflow_id, lambda run: run.state == RunState.PARKED) notification = await _wait_notification(rt, workflow_id, "delivered") @@ -731,4 +724,4 @@ async def test_concurrent_responds_resolve_to_one_answer(rt, deliver_spy): assert results.count("ok") == 1 assert _dbos_replies(rt, workflow_id) == 1 await _wait_run(rt, workflow_id, lambda run: run.state == RunState.FINISHED) - assert _notifications_for_run(rt, workflow_id)[0].state == "acknowledged" + assert (await _notifications_for_run(rt, workflow_id))[0].state == "acknowledged" diff --git a/backend/tests/test_oauth_client.py b/backend/tests/test_oauth_client.py index 8d8021c7..a19f3191 100644 --- a/backend/tests/test_oauth_client.py +++ b/backend/tests/test_oauth_client.py @@ -53,8 +53,10 @@ def _client(**overrides) -> OauthClient: return OauthClient(**kwargs) -def _connection(refresh_token: str = "rt-old", scopes: list[str] | None = None) -> OauthConnection: - return OauthConnection.create( +async def _connection( + refresh_token: str = "rt-old", scopes: list[str] | None = None +) -> OauthConnection: + return await OauthConnection.create( provider=_PROVIDER, account_id=SYSTEM_ACCOUNT_ID, refresh_token=refresh_token, @@ -75,7 +77,7 @@ def _lock_key(connection: OauthConnection) -> str: async def test_get_serves_the_cache_without_a_refresh(token_endpoint): - connection = _connection() + connection = await _connection() await get_client().set(_token_key(connection), "at-cached") token = await _client().get_access_token(connection=connection) @@ -86,13 +88,13 @@ async def test_get_serves_the_cache_without_a_refresh(token_endpoint): async def test_get_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_endpoint): token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} - connection = _connection() + connection = await _connection() token = await _client().get_access_token(connection=connection) assert token == "at-2" - db_session().expire_all() - assert OauthConnection.get(connection.id).refresh_token.decrypt() == "rt-new" + db_session().expunge_all() + assert (await OauthConnection.get(connection.id)).refresh_token.decrypt() == "rt-new" refresh = token_endpoint.requests[0] assert refresh["grant_type"] == "refresh_token" assert refresh["refresh_token"] == "rt-old" @@ -107,7 +109,7 @@ async def test_get_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_e async def test_get_fills_the_cache_only_after_the_rotation_is_saved(token_endpoint, monkeypatch): token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} - connection = _connection() + connection = await _connection() def _unsavable(self, rotated: str) -> None: raise RuntimeError("rotation write failed") @@ -122,7 +124,7 @@ def _unsavable(self, rotated: str) -> None: async def test_get_losing_the_lock_polls_for_the_winners_token(token_endpoint): - connection = _connection() + connection = await _connection() redis = get_client() await redis.set(_lock_key(connection), "1") @@ -139,7 +141,7 @@ async def _winner_finishes(): async def test_get_times_out_loudly_when_the_lock_never_frees(token_endpoint): - connection = _connection() + connection = await _connection() await get_client().set(_lock_key(connection), "1") with pytest.raises(OauthRefreshError, match="concurrent refresh"): @@ -148,19 +150,19 @@ async def test_get_times_out_loudly_when_the_lock_never_frees(token_endpoint): async def test_get_refresh_rejection_evicts_and_raises(token_endpoint): token_endpoint.status = 400 - connection = _connection() + connection = await _connection() with pytest.raises(OauthRefreshError, match="HTTP 400"): await _client().get_access_token(connection=connection) - assert OauthConnection.get(connection.id).refresh_token.decrypt() == "rt-old" + assert (await OauthConnection.get(connection.id)).refresh_token.decrypt() == "rt-old" redis = get_client() assert not await redis.get(_token_key(connection)) assert not await redis.get(_lock_key(connection)) async def test_get_refresh_uses_basic_auth(token_endpoint): - connection = _connection() + connection = await _connection() await _client(basic_auth=True).get_access_token(connection=connection) @@ -171,12 +173,12 @@ async def test_get_refresh_uses_basic_auth(token_endpoint): async def test_disconnect_revokes_the_connection_and_drops_the_cached_token(token_endpoint): - connection = _connection() + connection = await _connection() await get_client().set(_token_key(connection), "at-cached") await _client().disconnect(connection, reason="user") - revoked = OauthConnection.get(connection.id) + revoked = await OauthConnection.get(connection.id) assert revoked.revoked_at assert revoked.revoked_reason == "user" # Nothing secret outlives the consent at rest. @@ -186,17 +188,17 @@ async def test_disconnect_revokes_the_connection_and_drops_the_cached_token(toke # A second revoke keeps the first stamp. first_stamp = revoked.revoked_at - revoked.revoke("client_replaced") + await revoked.revoke("client_replaced") assert revoked.revoked_at == first_stamp assert revoked.revoked_reason == "user" async def test_a_revoke_landing_mid_refresh_is_not_overwritten(token_endpoint): - connection = _connection() + connection = await _connection() exchange = token_endpoint.handler - def revoke_then_rotate(request: httpx.Request) -> httpx.Response: - connection.revoke("user") + async def revoke_then_rotate(request: httpx.Request) -> httpx.Response: + await connection.revoke("user") return exchange(request) token_endpoint.handler = revoke_then_rotate @@ -205,13 +207,13 @@ def revoke_then_rotate(request: httpx.Request) -> httpx.Response: await _client().get_access_token(connection=connection) # The rotated token is not stored and no access token is cached. - assert not OauthConnection.get(connection.id).refresh_token + assert not (await OauthConnection.get(connection.id)).refresh_token assert not await get_client().get(_token_key(connection)) async def test_get_refuses_a_revoked_connection(token_endpoint): - connection = _connection() - connection.revoke("user") + connection = await _connection() + await connection.revoke("user") with pytest.raises(OauthRefreshError, match="revoked"): await _client().get_access_token(connection=connection) @@ -273,7 +275,7 @@ async def test_complete_connect_requires_a_refresh_token(token_endpoint): async def test_downscoped_get_asks_and_caches_apart_from_the_full_grant(token_endpoint): - connection = _connection(scopes=["posts.write", "profile.read"]) + connection = await _connection(scopes=["posts.write", "profile.read"]) token_endpoint.response = { "access_token": "at-narrow", "refresh_token": "rt-1", @@ -311,7 +313,7 @@ async def test_downscoped_get_asks_and_caches_apart_from_the_full_grant(token_en async def test_downscoped_get_rejects_scopes_outside_the_grant(token_endpoint): - connection = _connection(scopes=["profile.read"]) + connection = await _connection(scopes=["profile.read"]) with pytest.raises(OauthRefreshError, match="does not grant scope"): await _client().get_access_token(connection=connection, scopes=("posts.write",)) @@ -320,7 +322,7 @@ async def test_downscoped_get_rejects_scopes_outside_the_grant(token_endpoint): async def test_downscoped_get_rejects_a_provider_that_ignores_the_ask(token_endpoint): - connection = _connection(scopes=["posts.write", "profile.read"]) + connection = await _connection(scopes=["posts.write", "profile.read"]) token_endpoint.response = { "access_token": "at-broad", "refresh_token": "rt-1", @@ -336,7 +338,7 @@ async def test_downscoped_get_rejects_a_provider_that_ignores_the_ask(token_endp async def test_uncached_get_refreshes_past_a_live_cache_and_refills_it(token_endpoint): - connection = _connection() + connection = await _connection() redis = get_client() await redis.set(_token_key(connection), "at-tail") @@ -349,7 +351,7 @@ async def test_uncached_get_refreshes_past_a_live_cache_and_refills_it(token_end async def test_refresher_election_is_per_scope_set(token_endpoint): - connection = _connection(scopes=["profile.read"]) + connection = await _connection(scopes=["profile.read"]) redis = get_client() await redis.set(_lock_key(connection) + _scoped_suffix(("profile.read",)), "1") @@ -359,7 +361,7 @@ async def test_refresher_election_is_per_scope_set(token_endpoint): async def test_disconnect_evicts_the_scope_variant_keys(token_endpoint): - connection = _connection(scopes=["profile.read"]) + connection = await _connection(scopes=["profile.read"]) redis = get_client() scoped_key = _token_key(connection) + _scoped_suffix(("profile.read",)) await redis.set(_token_key(connection), "at-full") diff --git a/backend/tests/test_resume.py b/backend/tests/test_resume.py index 0e024063..8b7f648c 100644 --- a/backend/tests/test_resume.py +++ b/backend/tests/test_resume.py @@ -14,7 +14,7 @@ } -def _park(druks_db, *, context: str | None = None) -> None: +async def _park(druks_db, *, context: str | None = None) -> None: ask = dict(_ASK) if context is not None: ask["context"] = context @@ -27,8 +27,8 @@ def _park(druks_db, *, context: str | None = None) -> None: input_requested_at=Run.utc_now(), ) ) - druks_db.flush() - seed_dbos_status(druks_db, "r1", "parked") + await druks_db.flush() + await seed_dbos_status(druks_db, "r1", "parked") async def test_resume_sends_the_offered_control_as_the_action(druks_db, monkeypatch): @@ -38,7 +38,7 @@ async def fake_resume(self, **fields): captured.update(id=self.id, **fields) monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) await resume_run("r1", ResumeRequest(control="approve", answers={"q1": "a"})) assert captured == {"id": "r1", "action": "approve", "answers": {"q1": "a"}, "note": ""} @@ -54,7 +54,7 @@ async def fake_resume(self, **fields): captured.update(id=self.id, **fields) monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) await resume_run( "r1", @@ -79,7 +79,7 @@ async def fake_resume(self, **fields): raise AssertionError("must not resume on a rejected control") monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) with pytest.raises(HTTPException) as exc: await resume_run("r1", ResumeRequest(control="definitely-not-a-control")) @@ -92,7 +92,7 @@ async def fake_resume(self, **fields): raise AssertionError("must not resume on an invalid answer") monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) with pytest.raises(HTTPException) as exc: await resume_run("r1", ResumeRequest(control="approve", answers={"q9": "whatever"})) @@ -107,7 +107,7 @@ async def fake_resume(self, **fields): raise AssertionError("must not resume a guidance-free request_changes") monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db, context=context) + await _park(druks_db, context=context) with pytest.raises(HTTPException) as exc: await resume_run("r1", ResumeRequest(control="request_changes", note=" ")) @@ -122,7 +122,7 @@ async def fake_resume(self, **fields): captured.update(id=self.id, **fields) monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db, context="name the rollback boundary") + await _park(druks_db, context="name the rollback boundary") await resume_run("r1", ResumeRequest(control="request_changes", answers={}, note="")) @@ -141,7 +141,7 @@ async def fake_resume(self, **fields): captured.update(id=self.id, **fields) monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) await resume_run("r1", ResumeRequest(control="request_changes", note="split the migration")) assert captured == { @@ -157,7 +157,7 @@ async def fake_resume(self, **fields): raise AssertionError("must not resume on a blank answer") monkeypatch.setattr(Run, "resume", fake_resume) - _park(druks_db) + await _park(druks_db) with pytest.raises(HTTPException) as exc: await resume_run("r1", ResumeRequest(control="approve", answers={"q1": " "})) @@ -172,8 +172,8 @@ async def test_resume_404_when_run_missing(druks_db): async def test_resume_409_when_run_not_parked(druks_db): druks_db.add(Run(id="r2", kind="build")) - druks_db.flush() - seed_dbos_status(druks_db, "r2", "running") + await druks_db.flush() + await seed_dbos_status(druks_db, "r2", "running") with pytest.raises(HTTPException) as exc: await resume_run("r2", ResumeRequest(control="approve")) assert exc.value.status_code == 409 diff --git a/backend/tests/test_review.py b/backend/tests/test_review.py index f0d86ef2..755b15b0 100644 --- a/backend/tests/test_review.py +++ b/backend/tests/test_review.py @@ -38,23 +38,23 @@ def test_a_pull_requests_identity_is_its_handle(): @pytest.mark.parametrize("subject_id", ["acme/app", "acme/app#", "acme/app#0", "app#7", "#7"]) -def test_an_id_that_names_no_pull_request_is_a_miss(subject_id): - assert PullRequest.get_for_subject_id(subject_id) is None +async def test_an_id_that_names_no_pull_request_is_a_miss(subject_id): + assert await PullRequest.get_for_subject_id(subject_id) is None -def test_a_pull_request_heads_its_own_page(): - summary = PullRequest.get_for_subject_id("acme/app#7").get_summary() +async def test_a_pull_request_heads_its_own_page(): + summary = (await PullRequest.get_for_subject_id("acme/app#7")).get_summary() assert summary.repo == "acme/app" assert summary.pr_number == 7 assert summary.pull_request_url == "https://github.com/acme/app/pull/7" -def test_the_pull_request_board_and_page_mount(client: TestClient, druks_db): +async def test_the_pull_request_board_and_page_mount(client: TestClient, druks_db): # PullRequestReview declares PullRequest, so the app mounts its board and # page — keyed by a handle that carries both a path separator and a `#`. pull_request = PullRequest.get("acme/app", 7) - seed_run(druks_db, kind=PullRequestReview.kind, subject=pull_request, state="running") + await seed_run(druks_db, kind=PullRequestReview.kind, subject=pull_request, state="running") (row,) = client.get("/api/review/pull_request").json()["rows"] assert row["summary"]["id"] == "acme/app#7" @@ -72,7 +72,7 @@ def test_the_run_carries_the_pull_request_once(): assert list(PullRequestReview._run_input_model.model_fields) == ["requested_by"] -def test_a_queued_run_replays_through_its_subject(): +async def test_a_queued_run_replays_through_its_subject(): # A review enqueued before the repo and number came off the input still carries # them in its durable payload. The extra keys are ignored and the subject rides # separately, so the body binds and reads the pull request off the declaration. @@ -83,7 +83,8 @@ def test_a_queued_run_replays_through_its_subject(): ) assert run_kwargs == {"requested_by": "dev@example.com"} - assert (instance.subject.repo, instance.subject.number) == ("acme/app", 7) + subject = await instance.subject + assert (subject.repo, subject.number) == ("acme/app", 7) async def test_the_reviewer_prompt_names_the_pull_request_it_is_about(): @@ -130,44 +131,44 @@ async def test_comment_mode_reviews_publish_as_comments(): assert "`COMMENT` event" in output -def _connect_operator() -> None: - ServiceIdentity.connect( +async def _connect_operator() -> None: + await ServiceIdentity.connect( "github", identity={"app_id": "1", "slug": "druks-operator"}, secrets={"private_key": "operator-pem", "webhook_secret": "hook-secret"}, ) -def _set_review_setting(field: str, value: str) -> None: - SettingsOverride.set_app_setting("review", field, value, is_secret=True) +async def _set_review_setting(field: str, value: str) -> None: + await SettingsOverride.set_app_setting("review", field, value, is_secret=True) -def test_a_configured_review_identity_approves(druks_db): - _set_review_setting("app_id", "2") - _set_review_setting("private_key", "review-pem\nline-two") +async def test_a_configured_review_identity_approves(druks_db): + await _set_review_setting("app_id", "2") + await _set_review_setting("private_key", "review-pem\nline-two") - actor = get_review_actor() + actor = await get_review_actor() assert actor.mode == "approve" assert actor.client._app_id == "2" -def test_an_unset_review_identity_borrows_the_operator_in_comment_mode(druks_db): - _connect_operator() +async def test_an_unset_review_identity_borrows_the_operator_in_comment_mode(druks_db): + await _connect_operator() - actor = get_review_actor() + actor = await get_review_actor() assert actor.mode == "comment" assert actor.client._app_id == "1" -def test_a_half_configured_review_identity_still_borrows_the_operator(druks_db): +async def test_a_half_configured_review_identity_still_borrows_the_operator(druks_db): # Only a complete pair selects the distinct client; app_id alone is the # incoherent state clean() flags, not a mode switch. - _connect_operator() - _set_review_setting("app_id", "2") + await _connect_operator() + await _set_review_setting("app_id", "2") - actor = get_review_actor() + actor = await get_review_actor() assert actor.mode == "comment" assert actor.client._app_id == "1" @@ -192,14 +193,14 @@ def test_the_review_pem_declares_the_multiline_secret_presentation(): assert not field_multiline(Review.Settings.model_fields["app_id"]) -def test_review_identity_check_is_healthy_set_or_unset(druks_db): - assert check_review_identity().ok - assert "unset" in check_review_identity().detail +async def test_review_identity_check_is_healthy_set_or_unset(druks_db): + assert (await check_review_identity()).ok + assert "unset" in (await check_review_identity()).detail - _set_review_setting("app_id", "2") - _set_review_setting("private_key", "review-pem") + await _set_review_setting("app_id", "2") + await _set_review_setting("private_key", "review-pem") - result = check_review_identity() + result = await check_review_identity() assert result.ok assert "distinct App" in result.detail @@ -222,7 +223,7 @@ async def _start(cls, **kwargs): async def test_review_dispatch_starts_once_github_is_connected(druks_db, monkeypatch): - _connect_operator() + await _connect_operator() started = [] async def _start(cls, **kwargs): diff --git a/backend/tests/test_run_cancel.py b/backend/tests/test_run_cancel.py index 89e4038c..42ae400b 100644 --- a/backend/tests/test_run_cancel.py +++ b/backend/tests/test_run_cancel.py @@ -17,15 +17,15 @@ async def test_cancel_frees_subject_immediately(druks_db, monkeypatch): async def _dbos_cancel(workflow_id: str) -> None: cancelled.append(workflow_id) - druks_db.execute( + await druks_db.execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == workflow_id) .values(status="CANCELLED") ) monkeypatch.setattr("dbos.DBOS.cancel_workflow_async", _dbos_cancel) - note = Note.create(body="cancelled while parked") - run = seed_run( + note = await Note.create(body="cancelled while parked") + run = await seed_run( druks_db, kind=Summarize.kind, subject=note, @@ -39,8 +39,8 @@ async def _dbos_cancel(workflow_id: str) -> None: # cancel() never writes state — the already-loaded Run still carries the old # one until expired/re-selected, which is exactly what responses must do. # (cancel flushes the ambient session; the fixture session holds `run`.) - druks_db.flush() - druks_db.expire(run) + await druks_db.flush() + await druks_db.refresh(run) assert run.state == RunState.CANCELLED.value assert run.input_gate is None assert run.input_request is None diff --git a/backend/tests/test_run_state.py b/backend/tests/test_run_state.py index 7206870e..17eb0813 100644 --- a/backend/tests/test_run_state.py +++ b/backend/tests/test_run_state.py @@ -18,64 +18,64 @@ from uuid_utils import uuid7 -def _item_and_run(druks_db, state, **kwargs): - note = Note.create(body=f"run in {state}") - return note, seed_run(druks_db, kind=Summarize.kind, subject=note, state=state, **kwargs) +async def _item_and_run(druks_db, state, **kwargs): + note = await Note.create(body=f"run in {state}") + return note, await seed_run(druks_db, kind=Summarize.kind, subject=note, state=state, **kwargs) -def test_session_get_derives_state(druks_db): - _, run = _item_and_run(druks_db, "finished") - druks_db.expire_all() - assert Run.get(run.id).state == RunState.FINISHED.value +async def test_session_get_derives_state(druks_db): + _, run = await _item_and_run(druks_db, "finished") + druks_db.expunge_all() + assert (await Run.get(run.id)).state == RunState.FINISHED.value -def test_pending_splits_on_the_gate(druks_db): +async def test_pending_splits_on_the_gate(druks_db): # DBOS says PENDING either way; the gate is the one fact it can't know. - _, parked = _item_and_run(druks_db, "parked", input_gate="review_work") - _, live = _item_and_run(druks_db, "running") - druks_db.expire_all() - assert Run.get(parked.id).state == RunState.PARKED.value - assert Run.get(live.id).state == RunState.RUNNING.value + _, parked = await _item_and_run(druks_db, "parked", input_gate="review_work") + _, live = await _item_and_run(druks_db, "running") + druks_db.expunge_all() + assert (await Run.get(parked.id)).state == RunState.PARKED.value + assert (await Run.get(live.id)).state == RunState.RUNNING.value -def _rowless_run(session): +async def _rowless_run(session): """A run with no ``dbos.workflow_status`` row — the gap these tests are about, which ``seed_run`` closes by design.""" run = Run(id=str(uuid7()), kind=Summarize.kind, account_id="system") session.add(run) - session.flush() + await session.flush() return run -def test_fresh_run_without_a_dbos_row_reads_scheduled(druks_db): +async def test_fresh_run_without_a_dbos_row_reads_scheduled(druks_db): # start() writes the row before DBOS commits the enqueue; inside that gap a # brand-new run legitimately has no workflow_status row and reads scheduled. - run = _rowless_run(druks_db) - druks_db.expire_all() - assert Run.get(run.id).state == RunState.SCHEDULED.value + run = await _rowless_run(druks_db) + druks_db.expunge_all() + assert (await Run.get(run.id)).state == RunState.SCHEDULED.value -def test_run_without_a_dbos_row_past_grace_reads_orphaned(druks_db): +async def test_run_without_a_dbos_row_past_grace_reads_orphaned(druks_db): # A run still rowless past the grace window won't start — its DBOS row is # gone (system tables wiped, or the executor destroyed) — so derived state # reads orphaned instead of scheduled forever. - run = _rowless_run(druks_db) + run = await _rowless_run(druks_db) run.created_at = Base.utc_now() - timedelta(minutes=10) - druks_db.flush() - druks_db.expire_all() - assert Run.get(run.id).state == RunState.ORPHANED.value + await druks_db.flush() + druks_db.expunge_all() + assert (await Run.get(run.id)).state == RunState.ORPHANED.value -def test_unknown_dbos_status_reads_running(druks_db): +async def test_unknown_dbos_status_reads_running(druks_db): # A DBOS status this mapping predates must not crash reads. - _, run = _item_and_run(druks_db, "running") - druks_db.execute( + _, run = await _item_and_run(druks_db, "running") + await druks_db.execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == run.id) .values(status="SOME_FUTURE_STATUS") ) - druks_db.expire_all() - assert Run.get(run.id).state == RunState.RUNNING.value + druks_db.expunge_all() + assert (await Run.get(run.id)).state == RunState.RUNNING.value @pytest.mark.parametrize( @@ -86,22 +86,22 @@ def test_unknown_dbos_status_reads_running(druks_db): ("MAX_RECOVERY_ATTEMPTS_EXCEEDED", RunState.FAILED), ], ) -def test_statuses_the_seed_map_never_writes(druks_db, status, state): - _, run = _item_and_run(druks_db, "running") - druks_db.execute( +async def test_statuses_the_seed_map_never_writes(druks_db, status, state): + _, run = await _item_and_run(druks_db, "running") + await druks_db.execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == run.id) .values(status=status) ) - druks_db.expire_all() - assert Run.get(run.id).state == state.value + druks_db.expunge_all() + assert (await Run.get(run.id)).state == state.value -def test_queries_filter_on_derived_state(druks_db): - _, parked = _item_and_run(druks_db, "parked", input_gate="review_work") - _, done = _item_and_run(druks_db, "finished") +async def test_queries_filter_on_derived_state(druks_db): + _, parked = await _item_and_run(druks_db, "parked", input_gate="review_work") + _, done = await _item_and_run(druks_db, "finished") ids = set( - druks_db.scalars( + await druks_db.scalars( select(Run.id).where( Run.id.in_([parked.id, done.id]), Run.state.in_([RunState.PARKED.value, RunState.RUNNING.value]), @@ -111,18 +111,18 @@ def test_queries_filter_on_derived_state(druks_db): assert ids == {parked.id} -def test_updated_at_folds_in_the_dbos_write(druks_db): +async def test_updated_at_folds_in_the_dbos_write(druks_db): # DBOS stamps its updated_at in epoch milliseconds; the derived updated_at # converts it and wins over creation and the parked ask. - _, run = _item_and_run(druks_db, "finished") + _, run = await _item_and_run(druks_db, "finished") later_ms = int(datetime(2031, 1, 2, 3, 4, 5, tzinfo=UTC).timestamp() * 1000) - druks_db.execute( + await druks_db.execute( update(workflow_status) .where(workflow_status.c.workflow_uuid == run.id) .values(updated_at=later_ms) ) - druks_db.expire_all() - row = Run.get(run.id) + druks_db.expunge_all() + row = await Run.get(run.id) assert row.updated_at == datetime(2031, 1, 2, 3, 4, 5, tzinfo=UTC) assert row.updated_at > row.created_at @@ -143,7 +143,7 @@ async def test_facts_and_event_land_before_a_raising_subscriber(druks_db, _inlin # The fact write and its event commit before the signal fires, so a raising # subscriber can't roll them back. The failure itself still propagates: # delivery is at-least-once. - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") @subscribe(WorkflowEvent.PARKED, run=run.id) async def _raises(**_: object) -> None: @@ -157,14 +157,15 @@ async def _raises(**_: object) -> None: facts={"input_gate": "review_work", "input_request": {"label": "Review"}}, ) - ambient_session().expire_all() - row = Run.get(run.id) + ambient_session().expunge_all() + row = await Run.get(run.id) assert row.input_gate == "review_work" - events = ( - ambient_session() - .query(Event) - .filter_by(type="workflow.parked", subject_id=str(item.id)) - .all() + events = list( + ( + await ambient_session().execute( + select(Event).filter_by(type="workflow.parked", subject_id=str(item.id)) + ) + ).scalars() ) assert len(events) == 1 assert events[0].payload["gate"] == "review_work" @@ -176,12 +177,12 @@ async def test_lifecycle_subscribers_get_the_payload_before_dbos_commits(druks_d # derived state hasn't turned yet, which is why subscribers read the # payload, never Run.state. The body gets the run's own facts, never the # routing keys the filters match on. - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") seen: list[tuple[str, dict]] = [] @subscribe(WorkflowEvent.FINISHED, run=run.id) async def _reads_the_payload(**payload: object) -> None: - seen.append((Run.get(run.id).state, payload)) + seen.append(((await Run.get(run.id)).state, payload)) await _emit_run_event( run.id, @@ -202,7 +203,7 @@ async def test_cancellation_passes_through_untouched(druks_db, _inline_steps): # Operator cancel already carries its own reason and terminal status; the # body's cancellation exception must reach DBOS without a workflow.failed event # or a failure overwrite. - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") async def body() -> None: raise DBOSWorkflowCancelledError(f"workflow {run.id} cancelled") @@ -210,11 +211,12 @@ async def body() -> None: with pytest.raises(DBOSWorkflowCancelledError): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - assert Run.get(run.id).failure is None - types = [ - e.type for e in ambient_session().query(Event).filter_by(subject_id=str(item.id)).all() - ] + ambient_session().expunge_all() + assert (await Run.get(run.id)).failure is None + rows = ( + await ambient_session().execute(select(Event).filter_by(subject_id=str(item.id))) + ).scalars() + types = [e.type for e in rows] assert "workflow.failed" not in types @@ -226,7 +228,7 @@ async def test_failure_writes_the_reason_and_reraises(druks_db, _inline_steps): # derived state reads. from druks.durable.exceptions import FatalError - item, run = _item_and_run( + item, run = await _item_and_run( druks_db, "parked", input_gate="review_work", @@ -239,19 +241,18 @@ async def body() -> None: with pytest.raises(FatalError): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - row = Run.get(run.id) + ambient_session().expunge_all() + row = await Run.get(run.id) assert row.failure == "closed at review" # A bare FatalError carries no distinguishing code — only its message. assert row.failure_code == "" assert row.input_gate is None assert row.input_request is None failed = ( - ambient_session() - .query(Event) - .filter_by(type="workflow.failed", subject_id=str(item.id)) - .one() - ) + await ambient_session().execute( + select(Event).filter_by(type="workflow.failed", subject_id=str(item.id)) + ) + ).scalar_one() assert failed.payload["failure"] == "closed at review" @@ -261,7 +262,7 @@ async def test_gate_timeout_stamps_its_failure_code(druks_db, _inline_steps): # unanswered gate from a crash without parsing the failure text. from druks.durable.exceptions import GateTimeout - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") async def body() -> None: raise GateTimeout("review_work") @@ -269,15 +270,15 @@ async def body() -> None: with pytest.raises(GateTimeout): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - assert Run.get(run.id).failure_code == "gate_timeout" + ambient_session().expunge_all() + assert (await Run.get(run.id)).failure_code == "gate_timeout" @pytest.mark.asyncio async def test_a_harness_failure_stamps_its_code(druks_db, _inline_steps): from druks.harnesses.exceptions import HarnessOverloadedError - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") async def body() -> None: raise HarnessOverloadedError("claude exited with 1. API Error: 529 Overloaded.") @@ -285,8 +286,8 @@ async def body() -> None: with pytest.raises(HarnessOverloadedError): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - assert Run.get(run.id).failure_code == "overloaded" + ambient_session().expunge_all() + assert (await Run.get(run.id)).failure_code == "overloaded" @pytest.mark.asyncio @@ -296,7 +297,7 @@ async def test_an_exhausted_provisioning_failure_stamps_its_code(druks_db, _inli # SDK exception used to leave behind — so the dashboard/taxonomy can name it. from druks.harnesses.exceptions import HarnessSandboxProvisioningError - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") async def body() -> None: raise HarnessSandboxProvisioningError("exe.dev VM creation timed out") @@ -304,8 +305,8 @@ async def body() -> None: with pytest.raises(HarnessSandboxProvisioningError): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - assert Run.get(run.id).failure_code == "sandbox_provisioning" + ambient_session().expunge_all() + assert (await Run.get(run.id)).failure_code == "sandbox_provisioning" @pytest.mark.asyncio @@ -314,7 +315,7 @@ async def test_a_foreign_code_never_becomes_the_failure_code(druks_db, _inline_s the declaring families stamp the run; anything else records a crash.""" import asyncssh - item, run = _item_and_run(druks_db, "running") + item, run = await _item_and_run(druks_db, "running") async def body() -> None: raise asyncssh.PermissionDenied("denied") @@ -322,8 +323,8 @@ async def body() -> None: with pytest.raises(asyncssh.PermissionDenied): await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body) - ambient_session().expire_all() - assert Run.get(run.id).failure_code == "" + ambient_session().expunge_all() + assert (await Run.get(run.id)).failure_code == "" @pytest.mark.asyncio diff --git a/backend/tests/test_sandboxed_harness.py b/backend/tests/test_sandboxed_harness.py index e3aaca77..79e6574d 100644 --- a/backend/tests/test_sandboxed_harness.py +++ b/backend/tests/test_sandboxed_harness.py @@ -307,6 +307,10 @@ async def _download_artifacts(*args: Any, **kwargs: Any) -> None: assert downloads == [] +async def _harness_stub(model): + return ClaudeHarness + + async def test_run_prompt_builds_executes_and_parses( ctx: SimpleNamespace, ): @@ -329,7 +333,7 @@ class _FakeHarness: def mint_run_id(call_id: str | None) -> str: return call_id or "minted-id" - def build_invocation(self, **kwargs: Any) -> AgentInvocation: + async def build_invocation(self, **kwargs: Any) -> AgentInvocation: seen["build"] = kwargs return _inv(("claude", "--print")) @@ -491,15 +495,13 @@ def test_agent_result_names_the_agent_in_its_failure(): def _patch_harness_resolution(monkeypatch: pytest.MonkeyPatch) -> None: # run_agent resolves the harness + its settings from the DB; these tests # are about the failure boundary, not resolution. - from druks.harnesses.claude import ClaudeHarness - monkeypatch.setattr( - "druks.harnesses.registry.get_harness_for_model", lambda model: ClaudeHarness - ) - monkeypatch.setattr( - "druks.user_settings.models.HarnessSettings.require", - staticmethod(lambda name: SimpleNamespace(effort="high", timeout=60, fast_mode=False)), - ) + monkeypatch.setattr("druks.harnesses.registry.get_harness_for_model", _harness_stub) + + async def require(name): + return SimpleNamespace(effort="high", timeout=60, fast_mode=False) + + monkeypatch.setattr("druks.user_settings.models.HarnessSettings.require", staticmethod(require)) async def test_run_agent_carries_foreign_failures_as_harness_errors( diff --git a/backend/tests/test_secrets.py b/backend/tests/test_secrets.py index bdfe3650..59e5b5dd 100644 --- a/backend/tests/test_secrets.py +++ b/backend/tests/test_secrets.py @@ -4,6 +4,7 @@ import pytest from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.core.models import Uuid7Pk +from druks.database import db_session from druks.mcp.models import McpClientRegistration, McpServer from druks.models import Base from druks.secrets.exceptions import SecretDecryptError @@ -11,7 +12,7 @@ from druks.services.models import OauthConnection from druks.settings import load_settings from pydantic import ValidationError -from sqlalchemy import text +from sqlalchemy import select, text from sqlalchemy.exc import StatementError _TOKEN = "lin_secret_value" @@ -36,18 +37,20 @@ def _set_key(monkeypatch, tmp_path, value: str) -> None: monkeypatch.setenv("DRUKS_CONFIG", str(config_path)) -def _store_grant(refresh_token: str = "rt-secret", client_secret: str = "") -> OauthConnection: - server = McpServer.get_for_name("notion") or McpServer.create( +async def _store_grant( + refresh_token: str = "rt-secret", client_secret: str = "" +) -> OauthConnection: + server = await McpServer.get_for_name("notion") or await McpServer.create( name="notion", url="https://mcp.notion.test/sse" ) - McpClientRegistration.store( + await McpClientRegistration.store( server_id=server.id, account_id=SYSTEM_ACCOUNT_ID, token_endpoint="https://auth.test/token", client_id="client-123", client_secret=client_secret, ) - return OauthConnection.create( + return await OauthConnection.create( provider="mcp:notion", account_id=SYSTEM_ACCOUNT_ID, refresh_token=refresh_token, @@ -55,37 +58,37 @@ def _store_grant(refresh_token: str = "rt-secret", client_secret: str = "") -> O ) -def test_stored_secrets_are_ciphertext_and_reads_restore_them(druks_db): - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) +async def test_stored_secrets_are_ciphertext_and_reads_restore_them(druks_db): + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - blob = bytes(druks_db.execute(text("SELECT token FROM mcp_servers")).scalar_one()) + blob = bytes((await druks_db.execute(text("SELECT token FROM mcp_servers"))).scalar_one()) assert _TOKEN.encode() not in blob - druks_db.expire_all() - row = McpServer.get_for_name("linear") + druks_db.expunge_all() + row = await McpServer.get_for_name("linear") assert row.token.decrypt() == _TOKEN # The merged view every consumer reads carries the Secret itself, so the # plaintext exists only where decrypt() is called. - merged = McpServer._merged()["linear"] + merged = (await McpServer._merged())["linear"] assert merged["token"].decrypt() == _TOKEN -def test_grant_secret_halves_round_trip(druks_db): - _store_grant(refresh_token="rt-secret", client_secret="cs-secret") +async def test_grant_secret_halves_round_trip(druks_db): + await _store_grant(refresh_token="rt-secret", client_secret="cs-secret") - druks_db.expire_all() - grant = OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0] - registration = McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) + druks_db.expunge_all() + grant = (await OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID))[0] + registration = await McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) assert grant.refresh_token.decrypt() == "rt-secret" assert registration.client_secret.decrypt() == "cs-secret" -def test_loaded_secrets_are_lazy_and_redacted(monkeypatch, tmp_path, druks_db): - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - druks_db.expire_all() +async def test_loaded_secrets_are_lazy_and_redacted(monkeypatch, tmp_path, druks_db): + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) + druks_db.expunge_all() # Loading and logging a row never touches key material — decryption # happens only on decrypt(), and repr leaks nothing either way. - row = McpServer.get_for_name("linear") + row = await McpServer.get_for_name("linear") _set_key(monkeypatch, tmp_path, "") assert repr(row.token) == "Secret()" assert str(row.token) == "Secret()" @@ -93,25 +96,27 @@ def test_loaded_secrets_are_lazy_and_redacted(monkeypatch, tmp_path, druks_db): row.token.decrypt() -def test_empty_value_needs_no_key(monkeypatch, tmp_path, druks_db): +async def test_empty_value_needs_no_key(monkeypatch, tmp_path, druks_db): # "" stores as empty bytes — presence checks and decrypt() of an absent # secret never touch key material (proven by breaking the key first). - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token="") - druks_db.expire_all() + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token="") + druks_db.expunge_all() - assert bytes(druks_db.execute(text("SELECT token FROM mcp_servers")).scalar_one()) == b"" - row = McpServer.get_for_name("linear") + assert ( + bytes((await druks_db.execute(text("SELECT token FROM mcp_servers"))).scalar_one()) == b"" + ) + row = await McpServer.get_for_name("linear") _set_key(monkeypatch, tmp_path, "") assert not row.token assert row.token.decrypt() == "" -def test_non_str_assignment_is_rejected(druks_db): - server = McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) +async def test_non_str_assignment_is_rejected(druks_db): + server = await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) server.token = 123 with pytest.raises(StatementError, match="takes a str"): - druks_db.flush() + await db_session().flush() def test_missing_key_refuses_boot(monkeypatch, tmp_path): @@ -141,101 +146,102 @@ def test_malformed_key_refuses_boot(monkeypatch, tmp_path): load_settings() -def test_undecryptable_secret_raises_the_named_error(monkeypatch, tmp_path, druks_db): +async def test_undecryptable_secret_raises_the_named_error(monkeypatch, tmp_path, druks_db): # A key dropped from the list while rows written under it existed is the # usual cause — the error must say so, not surface a bare crypto traceback. - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - druks_db.expire_all() + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) + druks_db.expunge_all() _set_key(monkeypatch, tmp_path, _key()) with pytest.raises(SecretDecryptError, match="rotated out"): - McpServer.get_for_name("linear").token.decrypt() + (await McpServer.get_for_name("linear")).token.decrypt() -def test_garbled_envelope_raises_the_named_error(druks_db): +async def test_garbled_envelope_raises_the_named_error(druks_db): # No structural pre-checks in decrypt: GCM authentication (and the # ValueError a mangled nonce raises) fold every unreadable shape into the # one named error. - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - druks_db.execute(text(r"UPDATE mcp_servers SET token = '\x01ab'::bytea")) - druks_db.expire_all() + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) + await druks_db.execute(text(r"UPDATE mcp_servers SET token = '\x01ab'::bytea")) + druks_db.expunge_all() with pytest.raises(SecretDecryptError): - McpServer.get_for_name("linear").token.decrypt() + (await McpServer.get_for_name("linear")).token.decrypt() -def test_ciphertext_is_bound_to_its_column(druks_db): +async def test_ciphertext_is_bound_to_its_column(druks_db): # An envelope can't be replayed into any other encrypted column — not # another table's, and not a sibling column on the same row. - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - _store_grant(refresh_token="rt-secret", client_secret="cs-secret") - druks_db.execute( + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) + await _store_grant(refresh_token="rt-secret", client_secret="cs-secret") + await druks_db.execute( text( "UPDATE oauth_connections SET refresh_token =" " (SELECT token FROM mcp_servers WHERE name = 'linear')" ) ) - druks_db.execute( + await druks_db.execute( text( "UPDATE mcp_client_registrations SET client_secret =" " (SELECT refresh_token FROM oauth_connections WHERE provider = 'mcp:notion')" ) ) - druks_db.expire_all() + druks_db.expunge_all() - grant = OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0] - registration = McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) + grant = (await OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID))[0] + registration = await McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) with pytest.raises(SecretDecryptError): grant.refresh_token.decrypt() with pytest.raises(SecretDecryptError): registration.client_secret.decrypt() -def test_prepended_key_still_decrypts(monkeypatch, tmp_path, druks_db): +async def test_prepended_key_still_decrypts(monkeypatch, tmp_path, druks_db): # Rotation is prepend-only: new writes use the first key; rows written # under an older key keep decrypting as long as it stays in the list. old_key = _key() _set_key(monkeypatch, tmp_path, old_key) - McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) - _store_grant(refresh_token="rt-secret") + await McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) + await _store_grant(refresh_token="rt-secret") _set_key(monkeypatch, tmp_path, f"{_key()},{old_key}") - druks_db.expire_all() - assert McpServer.get_for_name("linear").token.decrypt() == _TOKEN - assert ( - OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0].refresh_token.decrypt() - == "rt-secret" - ) + druks_db.expunge_all() + assert (await McpServer.get_for_name("linear")).token.decrypt() == _TOKEN + assert (await OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID))[ + 0 + ].refresh_token.decrypt() == "rt-secret" # --- EncryptedJsonField (via the test-only model) --------------------------- -def test_json_mapping_round_trips_as_ciphertext(druks_db): +async def test_json_mapping_round_trips_as_ciphertext(druks_db): druks_db.add(EncryptedNote(data={"token": _TOKEN, "extra": "x"})) - druks_db.flush() + await druks_db.flush() - blob = bytes(druks_db.execute(text("SELECT data FROM test_encrypted_notes")).scalar_one()) + blob = bytes( + (await druks_db.execute(text("SELECT data FROM test_encrypted_notes"))).scalar_one() + ) assert _TOKEN.encode() not in blob - druks_db.expire_all() - note = druks_db.query(EncryptedNote).one() + druks_db.expunge_all() + note = (await druks_db.execute(select(EncryptedNote))).scalar_one() assert note.data["token"] == _TOKEN assert repr(note.data) == "SecretsMapping()" -def test_json_in_place_write_persists(druks_db): +async def test_json_in_place_write_persists(druks_db): # Writing one key of the mapping must mark the column dirty on its own # (the Mutable wiring) and survive the flush. druks_db.add(EncryptedNote(data={"token": "old"})) - druks_db.flush() - druks_db.expire_all() + await druks_db.flush() + druks_db.expunge_all() - note = druks_db.query(EncryptedNote).one() + note = (await druks_db.execute(select(EncryptedNote))).scalar_one() note.data["token"] = "new" - druks_db.flush() - druks_db.expire_all() + await druks_db.flush() + druks_db.expunge_all() - assert druks_db.query(EncryptedNote).one().data["token"] == "new" + assert (await druks_db.execute(select(EncryptedNote))).scalar_one().data["token"] == "new" def test_json_non_dict_assignment_is_rejected(druks_db): diff --git a/backend/tests/test_services.py b/backend/tests/test_services.py index 96060291..e486fcaf 100644 --- a/backend/tests/test_services.py +++ b/backend/tests/test_services.py @@ -19,70 +19,79 @@ from fastapi.testclient import TestClient from sqlalchemy import text + +def _server_app(): + from druks.api.server import app + + return app + + _PEM = "-----BEGIN RSA PRIVATE KEY-----\nline-one\nline-two\n-----END RSA PRIVATE KEY-----\n" _SECRET = "hook-secret-value" -def _connect( +async def _connect( *, app_id="12345", slug="druks-operator", private_key=_PEM, webhook_secret=_SECRET ) -> ServiceIdentity: - return ServiceIdentity.connect( + return await ServiceIdentity.connect( "github", identity={"app_id": app_id, "slug": slug}, secrets={"private_key": private_key, "webhook_secret": webhook_secret}, ) -def _github_entry(client: TestClient) -> dict: - entries = client.get("/api/services").json() +async def _github_entry(client) -> dict: + entries = (await client.get("/api/services")).json() return next(entry for entry in entries if entry["slug"] == "github") # --- The row ---------------------------------------------------------------- -def test_secrets_round_trip_and_rest_is_ciphertext(druks_db): - _connect() - druks_db.expire_all() +async def test_secrets_round_trip_and_rest_is_ciphertext(druks_db): + await _connect() + druks_db.expunge_all() - row = ServiceIdentity.get("github") + row = await ServiceIdentity.get("github") assert row.identity["app_id"] == "12345" assert row.identity["slug"] == "druks-operator" assert row.connected_at is not None assert row.secrets["private_key"] == _PEM assert row.secrets["webhook_secret"] == _SECRET - stored = druks_db.execute(text("SELECT secrets FROM service_identities")).scalar_one() + stored = (await druks_db.execute(text("SELECT secrets FROM service_identities"))).scalar_one() assert _PEM.encode() not in bytes(stored) assert _SECRET.encode() not in bytes(stored) -def test_connect_replaces_the_single_github_row(druks_db): - _connect() - _connect(app_id="777", slug="new-slug", private_key="new-pem", webhook_secret="new-secret") - druks_db.expire_all() +async def test_connect_replaces_the_single_github_row(druks_db): + await _connect() + await _connect( + app_id="777", slug="new-slug", private_key="new-pem", webhook_secret="new-secret" + ) + druks_db.expunge_all() - count = druks_db.execute(text("SELECT count(*) FROM service_identities")).scalar_one() + count = (await druks_db.execute(text("SELECT count(*) FROM service_identities"))).scalar_one() assert count == 1 - row = ServiceIdentity.get("github") + row = await ServiceIdentity.get("github") assert row.identity["app_id"] == "777" assert row.identity["slug"] == "new-slug" assert row.secrets["private_key"] == "new-pem" assert row.secrets["webhook_secret"] == "new-secret" -def test_get_raises_when_the_service_is_not_connected(druks_db): +async def test_get_raises_when_the_service_is_not_connected(druks_db): with pytest.raises(ServiceNotConnectedError, match="github is not connected"): - ServiceIdentity.get("github") + await ServiceIdentity.get("github") # --- The zero-argument client factory --------------------------------------- -def test_client_factory_resolves_only_the_row(druks_db): - _connect() +async def test_client_factory_resolves_only_the_row(druks_db): + await _connect() - client = get_github_client() + client = await get_github_client() assert client._app_id == "12345" assert client._private_key == _PEM @@ -91,14 +100,14 @@ def test_client_factory_resolves_only_the_row(druks_db): async def test_mention_handle_is_the_stored_slug(druks_db): # No transport stub: a refetch would ask GitHub and fail loudly here. - _connect() + await _connect() - assert await get_github_client().get_mention_handle() == "druks-operator" + assert await (await get_github_client()).get_mention_handle() == "druks-operator" -def test_client_factory_raises_the_typed_error_when_absent(druks_db): +async def test_client_factory_raises_the_typed_error_when_absent(druks_db): with pytest.raises(ServiceNotConnectedError): - get_github_client() + await get_github_client() # --- Webhook verification ---------------------------------------------------- @@ -117,26 +126,26 @@ def _events(body: bytes, signature: str | None, tmp_path) -> GitHubEvents: return events -def test_webhook_accepts_the_row_secrets_signature(druks_db, tmp_path): - _connect() +async def test_webhook_accepts_the_row_secrets_signature(druks_db, tmp_path): + await _connect() body = b'{"hello":"world"}' signature = "sha256=" + hmac.new(_SECRET.encode(), body, hashlib.sha256).hexdigest() - assert _events(body, signature, tmp_path).request_is_authentic() + assert await _events(body, signature, tmp_path).request_is_authentic() -def test_webhook_rejects_a_mismatched_signature(druks_db, tmp_path): - _connect() +async def test_webhook_rejects_a_mismatched_signature(druks_db, tmp_path): + await _connect() with pytest.raises(HTTPException) as raised: - _events(b"{}", "sha256=bogus", tmp_path).request_is_authentic() + await _events(b"{}", "sha256=bogus", tmp_path).request_is_authentic() assert raised.value.status_code == 401 -def test_webhook_rejects_a_missing_identity_before_dispatch(druks_db, tmp_path): +async def test_webhook_rejects_a_missing_identity_before_dispatch(druks_db, tmp_path): with pytest.raises(HTTPException) as raised: - _events(b"{}", "sha256=anything", tmp_path).request_is_authentic() + await _events(b"{}", "sha256=anything", tmp_path).request_is_authentic() assert raised.value.status_code == 401 assert "not connected" in raised.value.detail @@ -152,8 +161,8 @@ async def _slug(self) -> str: monkeypatch.setattr(GitHubClient, "get_authenticated_app_slug", _slug) -def test_list_reports_each_declared_service(druks_client: TestClient): - entry = _github_entry(druks_client) +async def test_list_reports_each_declared_service(druks_client: TestClient): + entry = await _github_entry(druks_client) assert entry["connected"] is False assert entry["required"] is True @@ -173,10 +182,12 @@ def test_list_reports_each_declared_service(druks_client: TestClient): assert [field["multiline"] for field in entry["fields"]] == [False, True, False] -def test_post_authenticates_then_creates_the_row(druks_client: TestClient, druks_db, monkeypatch): +async def test_post_authenticates_then_creates_the_row( + druks_client: TestClient, druks_db, monkeypatch +): _mock_authenticated_app(monkeypatch) - response = druks_client.post( + response = await druks_client.post( "/api/services/github", json={"app_id": "12345", "private_key": _PEM, "webhook_secret": _SECRET}, ) @@ -189,16 +200,16 @@ def test_post_authenticates_then_creates_the_row(druks_client: TestClient, druks assert _PEM not in response.text assert _SECRET not in response.text - connected = _github_entry(druks_client) + connected = await _github_entry(druks_client) assert connected["connected"] is True assert connected["facts"]["slug"] == "druks-operator" -def test_post_replaces_an_existing_row(druks_client: TestClient, druks_db, monkeypatch): - _connect() +async def test_post_replaces_an_existing_row(druks_client: TestClient, druks_db, monkeypatch): + await _connect() _mock_authenticated_app(monkeypatch, slug="replacement-app") - response = druks_client.post( + response = await druks_client.post( "/api/services/github", json={"app_id": "777", "private_key": "new-pem", "webhook_secret": "new-secret"}, ) @@ -207,23 +218,23 @@ def test_post_replaces_an_existing_row(druks_client: TestClient, druks_db, monke assert response.json()["facts"] == {"app_id": "777", "slug": "replacement-app"} -def test_post_rejects_an_unknown_service(druks_client: TestClient): - response = druks_client.post("/api/services/nope", json={"anything": "x"}) +async def test_post_rejects_an_unknown_service(druks_client: TestClient): + response = await druks_client.post("/api/services/nope", json={"anything": "x"}) assert response.status_code == 404 -def test_invalid_credentials_preserve_the_previous_row( +async def test_invalid_credentials_preserve_the_previous_row( druks_client: TestClient, druks_db, monkeypatch ): - _connect() + await _connect() async def _rejected(self) -> str: raise RuntimeError("boom-marker bad credentials") monkeypatch.setattr(GitHubClient, "get_authenticated_app_slug", _rejected) - response = druks_client.post( + response = await druks_client.post( "/api/services/github", json={"app_id": "999", "private_key": "bad-pem", "webhook_secret": "bad-secret"}, ) @@ -234,13 +245,13 @@ async def _rejected(self) -> str: assert "boom-marker" not in response.text assert "bad-pem" not in response.text - druks_db.expire_all() - row = ServiceIdentity.get("github") + druks_db.expunge_all() + row = await ServiceIdentity.get("github") assert row.identity["app_id"] == "12345" assert row.secrets["private_key"] == _PEM -def test_blank_fields_are_rejected_without_touching_github( +async def test_blank_fields_are_rejected_without_touching_github( druks_client: TestClient, druks_db, monkeypatch ): async def _never(self) -> str: @@ -248,14 +259,14 @@ async def _never(self) -> str: monkeypatch.setattr(GitHubClient, "get_authenticated_app_slug", _never) - response = druks_client.post( + response = await druks_client.post( "/api/services/github", json={"app_id": " ", "private_key": "", "webhook_secret": ""}, ) assert response.status_code == 422 with pytest.raises(ServiceNotConnectedError): - ServiceIdentity.get("github") + await ServiceIdentity.get("github") # --- The manifest flow (dashboard-created App) ------------------------------- @@ -266,12 +277,14 @@ def _manifest_from(page: str) -> dict: return json.loads(html.unescape(value)) -def test_manifest_page_submits_the_documented_app_to_github(druks_client: TestClient, tmp_path): - druks_client.app.state.settings = make_settings( +async def test_manifest_page_submits_the_documented_app_to_github( + druks_client: TestClient, tmp_path +): + _server_app().state.settings = make_settings( tmp_path, urls={"endpoint": "https://druks.example/"} ) - response = druks_client.get("/api/core/github/manifest") + response = await druks_client.get("/api/core/github/manifest") assert response.status_code == 200 assert 'action="https://github.com/settings/apps/new"' in response.text @@ -288,41 +301,43 @@ def test_manifest_page_submits_the_documented_app_to_github(druks_client: TestCl assert "pull_request_review" in manifest["default_events"] -def test_manifest_page_prefers_the_webhook_host_for_deliveries(druks_client: TestClient, tmp_path): - druks_client.app.state.settings = make_settings( +async def test_manifest_page_prefers_the_webhook_host_for_deliveries( + druks_client: TestClient, tmp_path +): + _server_app().state.settings = make_settings( tmp_path, urls={"endpoint": "https://druks.example", "webhook_host": "hooks.druks.example"}, ) - manifest = _manifest_from(druks_client.get("/api/core/github/manifest").text) + manifest = _manifest_from((await druks_client.get("/api/core/github/manifest")).text) assert ( manifest["hook_attributes"]["url"] == "https://hooks.druks.example/_external/github/events/" ) -def test_manifest_page_lets_the_operator_target_an_org(druks_client: TestClient, tmp_path): +async def test_manifest_page_lets_the_operator_target_an_org(druks_client: TestClient, tmp_path): # The org lives on the page, not in the card: naming one reroutes the # form to that org's create URL, and the input itself never reaches GitHub. - druks_client.app.state.settings = make_settings( + _server_app().state.settings = make_settings( tmp_path, urls={"endpoint": "https://druks.example"} ) - response = druks_client.get("/api/core/github/manifest") + response = await druks_client.get("/api/core/github/manifest") assert '' in response.text assert "https://github.com/organizations/" in response.text assert "this.elements.org.disabled = true" in response.text -def test_manifest_page_refuses_without_an_endpoint(druks_client: TestClient): - response = druks_client.get("/api/core/github/manifest") +async def test_manifest_page_refuses_without_an_endpoint(druks_client: TestClient): + response = await druks_client.get("/api/core/github/manifest") assert response.status_code == 409 assert "urls.endpoint" in response.json()["detail"] -def test_manifest_callback_exchanges_the_code_and_connects( +async def test_manifest_callback_exchanges_the_code_and_connects( druks_client: TestClient, druks_db, monkeypatch ): exchanged = [] @@ -342,7 +357,9 @@ async def fake_post(self, url, **kwargs): monkeypatch.setattr(httpx.AsyncClient, "post", fake_post) - response = druks_client.get("/api/core/github/manifest/callback", params={"code": "fresh-code"}) + response = await druks_client.get( + "/api/core/github/manifest/callback", params={"code": "fresh-code"} + ) assert response.status_code == 200 assert exchanged == ["https://api.github.com/app-manifests/fresh-code/conversions"] @@ -353,24 +370,28 @@ async def fake_post(self, url, **kwargs): # The page never carries the stored secrets. assert "line-one" not in response.text assert _SECRET not in response.text - druks_db.expire_all() - row = ServiceIdentity.get("github") + druks_db.expunge_all() + row = await ServiceIdentity.get("github") assert row.identity == {"app_id": "4242", "slug": "druks"} assert row.secrets == {"private_key": _PEM, "webhook_secret": _SECRET} -def test_manifest_callback_rejects_a_dead_code(druks_client: TestClient, druks_db, monkeypatch): +async def test_manifest_callback_rejects_a_dead_code( + druks_client: TestClient, druks_db, monkeypatch +): async def fake_post(self, url, **kwargs): return httpx.Response(404, json={"message": "Not Found"}) monkeypatch.setattr(httpx.AsyncClient, "post", fake_post) - response = druks_client.get("/api/core/github/manifest/callback", params={"code": "stale"}) + response = await druks_client.get( + "/api/core/github/manifest/callback", params={"code": "stale"} + ) assert response.status_code == 400 assert "restart" in response.json()["detail"] with pytest.raises(ServiceNotConnectedError): - ServiceIdentity.get("github") + await ServiceIdentity.get("github") # --- OAuth declaration -------------------------------------------------------- @@ -388,7 +409,7 @@ def declared_services(): services._items.update(saved) -def test_get_oauth_client_reads_the_connected_identity(declared_services, druks_db): +async def test_get_oauth_client_reads_the_connected_identity(declared_services, druks_db): from druks.services import Service from pydantic import BaseModel, SecretStr @@ -402,11 +423,11 @@ class Settings(BaseModel): client_id: str client_secret: SecretStr - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) - client = Acme.get_oauth_client() + client = await Acme.get_oauth_client() assert client.provider == "acme" assert client.authorization_endpoint == "https://acme.test/authorize" @@ -417,7 +438,7 @@ class Settings(BaseModel): assert client.extra_authorize_params == {"access_type": "offline"} -def test_oauth_service_declarations_fail_loudly(declared_services): +async def test_oauth_service_declarations_fail_loudly(declared_services): from druks.services import Service from pydantic import BaseModel, SecretStr @@ -452,7 +473,7 @@ class Settings(BaseModel): api_key: SecretStr with pytest.raises(TypeError, match="no OAuth endpoints"): - Plain.get_oauth_client() + await Plain.get_oauth_client() def test_abstract_base_shares_declarations_without_registering(declared_services): @@ -483,7 +504,7 @@ class Pinned(Service): slug = "pinned" -def test_with_scopes_declares_the_union_and_reads_connections(declared_services, monkeypatch): +async def test_with_scopes_declares_the_union_and_reads_connections(declared_services, monkeypatch): from druks.services import Service from pydantic import BaseModel, SecretStr @@ -515,26 +536,26 @@ class Digest: from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.services.models import OauthConnection - row = OauthConnection.create( + row = await OauthConnection.create( provider="acme", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt-1", scopes=["profile.read"], identity={"email": "night@acme.test"}, ) - connections = NightWatch.acme.list_for_account(SYSTEM_ACCOUNT_ID) + connections = await NightWatch.acme.list_for_account(SYSTEM_ACCOUNT_ID) assert [connection.id for connection in connections] == [row.id] assert connections[0].scopes == ["profile.read"] assert connections[0].identity == {"email": "night@acme.test"} assert connections[0].account_id == SYSTEM_ACCOUNT_ID - assert NightWatch.acme.get(row.id).id == row.id - assert not NightWatch.acme.get("missing") + assert (await NightWatch.acme.get(row.id)).id == row.id + assert not await NightWatch.acme.get("missing") # The handle serves live connections only; the revoked row survives. - row.revoke("user") - assert not NightWatch.acme.list_for_account(SYSTEM_ACCOUNT_ID) - assert not NightWatch.acme.get(row.id) - assert OauthConnection.get(row.id).identity == {"email": "night@acme.test"} + await row.revoke("user") + assert not await NightWatch.acme.list_for_account(SYSTEM_ACCOUNT_ID) + assert not await NightWatch.acme.get(row.id) + assert (await OauthConnection.get(row.id)).identity == {"email": "night@acme.test"} async def test_get_identity_without_a_declared_endpoint_is_empty(declared_services): @@ -637,14 +658,14 @@ def _complete_oauth_sign_in(client: TestClient, provider: str = "acme") -> None: assert response.status_code == 200 -def test_oauth_connect_redirects_to_consent_with_the_scope_union( +async def test_oauth_connect_redirects_to_consent_with_the_scope_union( tmp_path, acme, druks_db, monkeypatch ): from urllib.parse import parse_qsl, urlparse from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) @@ -706,7 +727,7 @@ async def test_oauth_callback_creates_and_reconnects_a_connection( from druks.services.models import OauthConnection from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) published = [] @@ -735,7 +756,7 @@ async def record(name, **kwargs): == 400 ) - [connection] = OauthConnection.list_for_provider("acme") + [connection] = await OauthConnection.list_for_provider("acme") assert connection.refresh_token.decrypt() == "rt-1" assert connection.identity == {"email": "op@acme.test"} assert connection.scopes == ["profile.read", "posts.write"] @@ -750,7 +771,7 @@ async def record(name, **kwargs): state = dict(parse_qsl(urlparse(reconnect.headers["location"]).query))["state"] finish = client.get("/api/oauth/callback", params={"state": state, "code": "c-2"}) assert finish.status_code == 200 - assert len(OauthConnection.list_for_provider("acme")) == 1 + assert len(await OauthConnection.list_for_provider("acme")) == 1 assert [name for name, _ in published] == ["oauth.connected", "oauth.connected"] fresh, reconsent = (kwargs for _, kwargs in published) @@ -766,10 +787,10 @@ async def record(name, **kwargs): assert not await get_client().get(stale_key) -def test_oauth_connect_rejects_an_unknown_reconnect_target(tmp_path, acme, druks_db): +async def test_oauth_connect_rejects_an_unknown_reconnect_target(tmp_path, acme, druks_db): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) @@ -777,16 +798,16 @@ def test_oauth_connect_rejects_an_unknown_reconnect_target(tmp_path, acme, druks assert client.get("/api/oauth/acme/connect?connection=zzz").status_code == 404 -def test_fresh_sign_in_with_matching_identity_resurrects_revoked_connection( +async def test_fresh_sign_in_with_matching_identity_resurrects_revoked_connection( tmp_path, keyed_acme, druks_db, oauth_events ): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( keyed_acme.slug, identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) - account = Account.get_or_create("op@example.com") - connection = OauthConnection.create( + account = await Account.get_or_create("op@example.com") + connection = await OauthConnection.create( provider=keyed_acme.slug, account_id=account.id, refresh_token="rt-old", @@ -794,20 +815,22 @@ def test_fresh_sign_in_with_matching_identity_resurrects_revoked_connection( identity={"sub": "account-1"}, ) connection_id = connection.id - connection.revoke("user") + await connection.revoke("user") settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) with TestClient(configure_app_for_test(settings=settings)) as client: _complete_oauth_sign_in(client, keyed_acme.slug) - db_session().expire_all() - resurrected = OauthConnection.get(connection_id) + db_session().expunge_all() + resurrected = await OauthConnection.get(connection_id) assert resurrected assert not resurrected.revoked_at assert not resurrected.revoked_reason assert resurrected.refresh_token.decrypt() == "rt-1" - assert [row.id for row in OauthConnection.list_for_provider(keyed_acme.slug)] == [connection_id] - assert len(OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 1 + assert [row.id for row in await OauthConnection.list_for_provider(keyed_acme.slug)] == [ + connection_id + ] + assert len(await OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 1 assert oauth_events == [ ( "oauth.connected", @@ -821,16 +844,16 @@ def test_fresh_sign_in_with_matching_identity_resurrects_revoked_connection( ] -def test_matching_fresh_sign_in_lands_on_live_connection_and_evicts_cached_token( +async def test_matching_fresh_sign_in_lands_on_live_connection_and_evicts_cached_token( tmp_path, keyed_acme, druks_db, oauth_events, monkeypatch ): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( keyed_acme.slug, identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) - account = Account.get_or_create("op@example.com") - connection = OauthConnection.create( + account = await Account.get_or_create("op@example.com") + connection = await OauthConnection.create( provider=keyed_acme.slug, account_id=account.id, refresh_token="rt-old", @@ -850,26 +873,26 @@ async def record_eviction(oauth_client, evicted_connection_id): with TestClient(configure_app_for_test(settings=settings)) as client: _complete_oauth_sign_in(client, keyed_acme.slug) - db_session().expire_all() - reconnected = OauthConnection.get(connection_id) + db_session().expunge_all() + reconnected = await OauthConnection.get(connection_id) assert reconnected assert reconnected.refresh_token.decrypt() == "rt-1" - assert len(OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 1 + assert len(await OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 1 assert evicted_connection_ids == [connection_id] assert oauth_events[-1][1]["connection_id"] == connection_id assert oauth_events[-1][1]["reconsent"] is True -def test_fresh_sign_in_with_live_and_revoked_identity_matches_lands_on_live_connection( +async def test_fresh_sign_in_with_live_and_revoked_identity_matches_lands_on_live_connection( tmp_path, keyed_acme, druks_db, oauth_events ): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( keyed_acme.slug, identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) - account = Account.get_or_create("op@example.com") - live = OauthConnection.create( + account = await Account.get_or_create("op@example.com") + live = await OauthConnection.create( provider=keyed_acme.slug, account_id=account.id, refresh_token="rt-live-old", @@ -877,7 +900,7 @@ def test_fresh_sign_in_with_live_and_revoked_identity_matches_lands_on_live_conn identity={"sub": "account-1"}, ) live_id = live.id - revoked = OauthConnection.create( + revoked = await OauthConnection.create( provider=keyed_acme.slug, account_id=account.id, refresh_token="rt-revoked-old", @@ -885,35 +908,35 @@ def test_fresh_sign_in_with_live_and_revoked_identity_matches_lands_on_live_conn identity={"sub": "account-1"}, ) revoked_id = revoked.id - revoked.revoke("user") + await revoked.revoke("user") settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) with TestClient(configure_app_for_test(settings=settings)) as client: _complete_oauth_sign_in(client, keyed_acme.slug) - db_session().expire_all() - reconnected = OauthConnection.get(live_id) - still_revoked = OauthConnection.get(revoked_id) + db_session().expunge_all() + reconnected = await OauthConnection.get(live_id) + still_revoked = await OauthConnection.get(revoked_id) assert reconnected assert still_revoked assert reconnected.refresh_token.decrypt() == "rt-1" assert still_revoked.revoked_at assert not still_revoked.refresh_token - assert len(OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 2 + assert len(await OauthConnection.list_for_provider(keyed_acme.slug, include_revoked=True)) == 2 assert oauth_events[-1][1]["connection_id"] == live_id assert oauth_events[-1][1]["reconsent"] is True -def test_fresh_sign_in_without_the_declared_identity_fact_creates_a_new_connection( +async def test_fresh_sign_in_without_the_declared_identity_fact_creates_a_new_connection( tmp_path, acme, druks_db, oauth_events ): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( acme.slug, identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) - account = Account.get_or_create("op@example.com") - revoked = OauthConnection.create( + account = await Account.get_or_create("op@example.com") + revoked = await OauthConnection.create( provider=acme.slug, account_id=account.id, refresh_token="rt-old", @@ -921,28 +944,30 @@ def test_fresh_sign_in_without_the_declared_identity_fact_creates_a_new_connecti identity={"sub": "account-1"}, ) revoked_id = revoked.id - revoked.revoke("user") + await revoked.revoke("user") settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) with TestClient(configure_app_for_test(settings=settings)) as client: _complete_oauth_sign_in(client, acme.slug) - db_session().expire_all() - [created] = OauthConnection.list_for_provider(acme.slug) + db_session().expunge_all() + [created] = await OauthConnection.list_for_provider(acme.slug) assert created.id != revoked_id - assert len(OauthConnection.list_for_provider(acme.slug, include_revoked=True)) == 2 - assert OauthConnection.get(revoked_id).revoked_at + assert len(await OauthConnection.list_for_provider(acme.slug, include_revoked=True)) == 2 + assert (await OauthConnection.get(revoked_id)).revoked_at assert oauth_events[-1][1]["connection_id"] == created.id assert oauth_events[-1][1]["reconsent"] is False -def test_fresh_sign_in_after_revoke_creates_a_new_connection(tmp_path, acme, druks_db, monkeypatch): +async def test_fresh_sign_in_after_revoke_creates_a_new_connection( + tmp_path, acme, druks_db, monkeypatch +): from urllib.parse import parse_qsl, urlparse from druks.services.models import OauthConnection from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) published = [] @@ -965,15 +990,15 @@ def sign_in(): ) sign_in() - [first] = OauthConnection.list_for_provider("acme") + [first] = await OauthConnection.list_for_provider("acme") assert client.delete(f"/api/oauth/connections/{first.id}").status_code == 204 # The identity facts do not include "sub", so the sign-in cannot # match the existing row. The revoked row stays as history. sign_in() - [live] = OauthConnection.list_for_provider("acme") + [live] = await OauthConnection.list_for_provider("acme") assert live.id != first.id - assert len(OauthConnection.list_for_provider("acme", include_revoked=True)) == 2 + assert len(await OauthConnection.list_for_provider("acme", include_revoked=True)) == 2 events = [name for name, _ in published] assert events == ["oauth.connected", "oauth.disconnected", "oauth.connected"] @@ -985,13 +1010,15 @@ def sign_in(): } -def test_reconsent_returns_a_revoked_connection_to_life(tmp_path, acme, druks_db, monkeypatch): +async def test_reconsent_returns_a_revoked_connection_to_life( + tmp_path, acme, druks_db, monkeypatch +): from urllib.parse import parse_qsl, urlparse from druks.services.models import OauthConnection from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) published = [] @@ -1005,7 +1032,7 @@ async def record(name, **kwargs): consent = client.get("/api/oauth/acme/connect", follow_redirects=False) state = dict(parse_qsl(urlparse(consent.headers["location"]).query))["state"] client.get("/api/oauth/callback", params={"state": state, "code": "c-1"}) - [connection] = OauthConnection.list_for_provider("acme") + [connection] = await OauthConnection.list_for_provider("acme") assert client.delete(f"/api/oauth/connections/{connection.id}").status_code == 204 # Reconsent names the row and makes the revoked consent live again. @@ -1019,8 +1046,8 @@ async def record(name, **kwargs): ) # The routes wrote in their own transactions; drop stale instances. - db_session().expire_all() - [live] = OauthConnection.list_for_provider("acme") + db_session().expunge_all() + [live] = await OauthConnection.list_for_provider("acme") assert live.id == connection.id assert not live.revoked_at assert not live.revoked_reason @@ -1034,7 +1061,7 @@ async def record(name, **kwargs): } -def test_connections_list_and_revoke(tmp_path, acme, druks_db, monkeypatch): +async def test_connections_list_and_revoke(tmp_path, acme, druks_db, monkeypatch): from druks.accounts.models import Account from druks.services.models import OauthConnection from druks.testing import configure_app_for_test @@ -1045,9 +1072,9 @@ async def record(name, **kwargs): published.append((name, kwargs)) monkeypatch.setattr("druks.services.routes.publish", record) - me = Account.get_or_create("op@example.com") + me = await Account.get_or_create("op@example.com") with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - row = OauthConnection.create( + row = await OauthConnection.create( provider="acme", account_id=me.id, refresh_token="rt-1", scopes=["profile.read"] ) [listed] = client.get("/api/oauth/connections").json() @@ -1057,10 +1084,10 @@ async def record(name, **kwargs): assert listed["revokedAt"] is None assert client.delete(f"/api/oauth/connections/{row.id}").status_code == 204 - assert not OauthConnection.list_for_provider("acme") + assert not await OauthConnection.list_for_provider("acme") # The route revoked in its own transaction; drop the stale instance. - db_session().expire_all() - revoked = OauthConnection.get(row.id) + db_session().expunge_all() + revoked = await OauthConnection.get(row.id) assert revoked.revoked_at assert revoked.revoked_reason == "user" assert not revoked.refresh_token @@ -1079,7 +1106,7 @@ async def record(name, **kwargs): ] -def test_replacing_the_client_credentials_revokes_its_connections( +async def test_replacing_the_client_credentials_revokes_its_connections( tmp_path, acme, druks_db, monkeypatch ): from druks.accounts.constants import SYSTEM_ACCOUNT_ID @@ -1092,7 +1119,7 @@ async def record(name, **kwargs): published.append((name, kwargs)) monkeypatch.setattr("druks.services.routes.publish", record) - row = OauthConnection.create( + row = await OauthConnection.create( provider="acme", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt-old", scopes=[] ) @@ -1103,9 +1130,9 @@ async def record(name, **kwargs): assert response.status_code == 200 # The new client can never refresh the old client's connections. - db_session().expire_all() - assert not OauthConnection.list_for_provider("acme") - [revoked] = OauthConnection.list_for_provider("acme", include_revoked=True) + db_session().expunge_all() + assert not await OauthConnection.list_for_provider("acme") + [revoked] = await OauthConnection.list_for_provider("acme", include_revoked=True) assert revoked.id == row.id assert revoked.revoked_reason == "client_replaced" assert not revoked.refresh_token @@ -1117,41 +1144,41 @@ async def record(name, **kwargs): ] -def test_list_serves_the_connections_beside_the_declared_union(tmp_path, acme, druks_db): +async def test_list_serves_the_connections_beside_the_declared_union(tmp_path, acme, druks_db): from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.services.models import OauthConnection from druks.testing import configure_app_for_test - def entry(client, slug="acme"): + async def entry(client, slug="acme"): return next(e for e in client.get("/api/services").json() if e["slug"] == slug) with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: - assert entry(client, "github")["isOauth"] is False - before = entry(client) + assert (await entry(client, "github"))["isOauth"] is False + before = await entry(client) assert before["isOauth"] is True assert before["connections"] == [] assert before["requiredScopes"] == ["openid", "posts.write", "profile.read"] assert before["usedBy"] == ["night_watch.acme"] - row = OauthConnection.create( + row = await OauthConnection.create( provider="acme", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt-1", scopes=["profile.read"], ) - [connection] = entry(client)["connections"] + [connection] = (await entry(client))["connections"] assert connection["id"] == row.id assert connection["scopes"] == ["profile.read"] assert connection["identity"] == {} assert connection["connectedAt"] -def test_next_lands_the_user_back_on_the_app_page(tmp_path, acme, druks_db): +async def test_next_lands_the_user_back_on_the_app_page(tmp_path, acme, druks_db): from urllib.parse import parse_qsl, urlparse from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) @@ -1170,10 +1197,10 @@ def test_next_lands_the_user_back_on_the_app_page(tmp_path, acme, druks_db): assert finish.headers["location"] == "/app/night_watch/accounts" -def test_next_rejects_anything_but_a_bare_path(tmp_path, acme, druks_db): +async def test_next_rejects_anything_but_a_bare_path(tmp_path, acme, druks_db): from druks.testing import configure_app_for_test - ServiceIdentity.connect( + await ServiceIdentity.connect( "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} ) settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) diff --git a/backend/tests/test_settings_overrides.py b/backend/tests/test_settings_overrides.py index c804359e..d4540b87 100644 --- a/backend/tests/test_settings_overrides.py +++ b/backend/tests/test_settings_overrides.py @@ -5,19 +5,19 @@ # primitives with no API-level twin. -def test_app_setting_override_then_default(druks_db): +async def test_app_setting_override_then_default(druks_db): # No override → the declared default passed by the caller. - assert SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is True + assert await SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is True # An override wins — including turning it off. - SettingsOverride.set_app_setting("ship", "auto_merge", False, is_secret=False) - assert SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is False + await SettingsOverride.set_app_setting("ship", "auto_merge", False, is_secret=False) + assert await SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is False # Clearing reverts to the caller's default. - SettingsOverride.set_app_setting("ship", "auto_merge", None, is_secret=False) - assert SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is True + await SettingsOverride.set_app_setting("ship", "auto_merge", None, is_secret=False) + assert await SettingsOverride.app_setting("ship", "auto_merge", True, is_secret=False) is True -def test_workflow_setting_namespaced_by_kind(druks_db): - SettingsOverride.set_workflow_setting("ship.build", "shared", "a") - SettingsOverride.set_workflow_setting("other_workflow", "shared", "b") - assert SettingsOverride.workflow_setting("ship.build", "shared", None) == "a" - assert SettingsOverride.workflow_setting("other_workflow", "shared", None) == "b" +async def test_workflow_setting_namespaced_by_kind(druks_db): + await SettingsOverride.set_workflow_setting("ship.build", "shared", "a") + await SettingsOverride.set_workflow_setting("other_workflow", "shared", "b") + assert await SettingsOverride.workflow_setting("ship.build", "shared", None) == "a" + assert await SettingsOverride.workflow_setting("other_workflow", "shared", None) == "b" diff --git a/backend/tests/test_signals.py b/backend/tests/test_signals.py index d3d6a128..79cae233 100644 --- a/backend/tests/test_signals.py +++ b/backend/tests/test_signals.py @@ -1,13 +1,13 @@ import pytest from druks.apps.exceptions import SubscriberDeclarationError +from druks.database import db_session from druks.signals import publish, subscribe from druks_field_notes.models import Note from druks_field_notes.workflows import Summarize -from sqlalchemy.orm import object_session -def _note(body: str = "probe") -> Note: - return Note.create(body=body) +async def _note(body: str = "probe") -> Note: + return await Note.create(body=body) @pytest.mark.asyncio @@ -25,7 +25,7 @@ async def boom(**_: object) -> None: @pytest.mark.asyncio async def test_subject_subscriber_receives_the_row(druks_db): - item = _note() + item = await _note() received = [] @subscribe("test.subject_row", subject=Note) @@ -39,7 +39,7 @@ async def receive(*, subject: Note) -> None: @pytest.mark.asyncio async def test_deleted_subject_skips_the_subscriber(druks_db): - item = _note() + item = await _note() identity = item.identity received = [] @@ -47,8 +47,8 @@ async def test_deleted_subject_skips_the_subscriber(druks_db): async def receive(*, subject: Note) -> None: received.append(subject) - object_session(item).delete(item) - object_session(item).flush() + await db_session().delete(item) + await db_session().flush() await publish("test.deleted_subject", subject=identity) assert received == [] @@ -56,7 +56,7 @@ async def receive(*, subject: Note) -> None: @pytest.mark.asyncio async def test_another_subjects_event_skips_the_subscriber(druks_db): - item = _note() + item = await _note() received = [] @subscribe("test.other_subject", subject=Note) @@ -74,7 +74,7 @@ async def test_workflow_filter_narrows_to_that_workflow(druks_db): # The body names the fact it came for; the kind it was matched on is routing, # so it never reaches the signature. Summarize declares its subject, so the # filter narrows to notes too and the body is handed the row. - item = _note() + item = await _note() received = [] @subscribe("test.workflow_filter", workflow=Summarize) @@ -113,9 +113,9 @@ def test_a_subscriber_asking_for_routing_fails_at_declaration(): async def receive(*, run: str, kind: str, **_: object) -> None: ... -def test_a_run_resolves_its_subject_through_the_declaration(druks_db): - item = _note() +async def test_a_run_resolves_its_subject_through_the_declaration(druks_db): + item = await _note() workflow = Summarize() workflow._subject = item.identity - assert workflow.subject is item + assert await workflow.subject is item diff --git a/backend/tests/test_skills.py b/backend/tests/test_skills.py index 64207250..e72a198c 100644 --- a/backend/tests/test_skills.py +++ b/backend/tests/test_skills.py @@ -4,6 +4,7 @@ from pathlib import Path import pytest +from druks.database import db_session from druks.skills import install as install_mod from druks.skills import routes as routes_mod from druks.skills.datastructures import CollectionContents, InstalledSkill @@ -104,8 +105,8 @@ def test_non_github_url_rejected(): install_mod._parse_github_repo("https://gitlab.com/owner/repo") -def test_collection_create_get_cascade_delete(druks_db): - collection = SkillCollection.create( +async def test_collection_create_get_cascade_delete(druks_db): + collection = await SkillCollection.create( source="https://github.com/o/r", name="o/r", skills=[ @@ -113,26 +114,26 @@ def test_collection_create_get_cascade_delete(druks_db): InstalledSkill(name="beta", description="two", path="/p/beta", content_hash="b"), ], ) - assert SkillCollection.get_for_source("https://github.com/o/r").id == collection.id - assert Skill.installed_names() == {"alpha", "beta"} - assert [c.name for c in SkillCollection.list_all()] == ["o/r"] + assert (await SkillCollection.get_for_source("https://github.com/o/r")).id == collection.id + assert await Skill.installed_names() == {"alpha", "beta"} + assert [c.name for c in await SkillCollection.list_all()] == ["o/r"] - assert [skill.name for skill in Skill.list_delivered(())] == ["alpha", "beta"] - assert Skill.delivery_excludes(()) == () - assert [skill.name for skill in Skill.list_delivered(("alpha",))] == ["alpha"] - assert Skill.delivery_excludes(("alpha",)) == ("./beta",) + assert [skill.name for skill in await Skill.list_delivered(())] == ["alpha", "beta"] + assert await Skill.delivery_excludes(()) == () + assert [skill.name for skill in await Skill.list_delivered(("alpha",))] == ["alpha"] + assert await Skill.delivery_excludes(("alpha",)) == ("./beta",) - Skill.get("alpha").enabled = False - druks_db.flush() + (await Skill.get("alpha")).enabled = False + await druks_db.flush() - assert [skill.name for skill in Skill.list_delivered(())] == ["beta"] - assert Skill.delivery_excludes(()) == ("./alpha",) - assert Skill.list_delivered(("alpha",)) == [] - assert Skill.delivery_excludes(("alpha",)) == ("./alpha", "./beta") + assert [skill.name for skill in await Skill.list_delivered(())] == ["beta"] + assert await Skill.delivery_excludes(()) == ("./alpha",) + assert await Skill.list_delivered(("alpha",)) == [] + assert await Skill.delivery_excludes(("alpha",)) == ("./alpha", "./beta") - collection.delete() - assert SkillCollection.list_all() == [] - assert Skill.installed_names() == set() + await collection.delete() + assert await SkillCollection.list_all() == [] + assert await Skill.installed_names() == set() def test_collection_routes_install_list_remove(tmp_path, monkeypatch): @@ -178,7 +179,7 @@ async def fake_fetch(url, skills_dir, reserved_names): assert client.get("/api/skills").json() == [] -def test_sync_updates_changed_skill_and_timestamps(tmp_path, monkeypatch, druks_db): +async def test_sync_updates_changed_skill_and_timestamps(tmp_path, monkeypatch, druks_db): settings = make_settings(tmp_path) original_skill = _skill_md("alpha") updated_skill = _skill_md("alpha", "updated description") @@ -190,12 +191,12 @@ def test_sync_updates_changed_skill_and_timestamps(tmp_path, monkeypatch, druks_ monkeypatch, {"alpha/SKILL.md": original_skill}, ) - collection = SkillCollection.get(collection_id) - skill = Skill.get("alpha") + collection = await SkillCollection.get(collection_id) + skill = await Skill.get("alpha") original_hash = skill.content_hash collection.updated_at = old_timestamp skill.updated_at = old_timestamp - druks_db.flush() + await db_session().flush() _patch_download( monkeypatch, @@ -205,7 +206,8 @@ def test_sync_updates_changed_skill_and_timestamps(tmp_path, monkeypatch, druks_ assert response.status_code == 200 assert response.json()["skills"][0]["description"] == "updated description" - druks_db.expire_all() + await db_session().refresh(skill) + await db_session().refresh(collection) assert skill.description == "updated description" assert skill.content_hash != original_hash assert skill.updated_at > old_timestamp @@ -214,16 +216,17 @@ def test_sync_updates_changed_skill_and_timestamps(tmp_path, monkeypatch, druks_ skill_timestamp = skill.updated_at collection.updated_at = old_timestamp - druks_db.flush() + await db_session().flush() response = client.post(f"/api/skills/{collection_id}/sync") assert response.status_code == 200 - druks_db.expire_all() + await db_session().refresh(skill) + await db_session().refresh(collection) assert skill.updated_at == skill_timestamp assert collection.updated_at > old_timestamp -def test_sync_adds_new_skill(tmp_path, monkeypatch, druks_db): +async def test_sync_adds_new_skill(tmp_path, monkeypatch, druks_db): settings = make_settings(tmp_path) alpha_skill = _skill_md("alpha") beta_skill = _skill_md("beta") @@ -249,12 +252,12 @@ def test_sync_adds_new_skill(tmp_path, monkeypatch, druks_db): assert response.status_code == 200 assert [skill["name"] for skill in response.json()["skills"]] == ["alpha", "beta"] - druks_db.expire_all() - assert Skill.get("beta").collection_id == collection_id + druks_db.expunge_all() + assert (await Skill.get("beta")).collection_id == collection_id assert (settings.skills_dir / "beta" / "SKILL.md").read_bytes() == beta_skill -def test_sync_removes_missing_skill_and_files(tmp_path, monkeypatch, druks_db): +async def test_sync_removes_missing_skill_and_files(tmp_path, monkeypatch, druks_db): settings = make_settings(tmp_path) alpha_skill = _skill_md("alpha") beta_skill = _skill_md("beta") @@ -279,12 +282,12 @@ def test_sync_removes_missing_skill_and_files(tmp_path, monkeypatch, druks_db): assert response.status_code == 200 assert [skill["name"] for skill in response.json()["skills"]] == ["alpha"] - druks_db.expire_all() - assert not Skill.get("beta") + druks_db.expunge_all() + assert not await Skill.get("beta") assert not beta_path.exists() -def test_sync_preserves_disabled_skill(tmp_path, monkeypatch, druks_db): +async def test_sync_preserves_disabled_skill(tmp_path, monkeypatch, druks_db): alpha_skill = _skill_md("alpha") with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: @@ -293,9 +296,9 @@ def test_sync_preserves_disabled_skill(tmp_path, monkeypatch, druks_db): monkeypatch, {"alpha/SKILL.md": alpha_skill}, ) - skill = Skill.get("alpha") + skill = await Skill.get("alpha") skill.enabled = False - druks_db.flush() + await db_session().flush() _patch_download( monkeypatch, _tarball("owner-repo-abc", {"alpha/SKILL.md": alpha_skill}), @@ -305,7 +308,7 @@ def test_sync_preserves_disabled_skill(tmp_path, monkeypatch, druks_db): assert response.status_code == 200 assert response.json()["skills"][0]["enabled"] is False - druks_db.expire_all() + await db_session().refresh(skill) assert skill.enabled is False diff --git a/backend/tests/test_spa_serving.py b/backend/tests/test_spa_serving.py index 0f0957f5..b8b28347 100644 --- a/backend/tests/test_spa_serving.py +++ b/backend/tests/test_spa_serving.py @@ -60,34 +60,30 @@ def test_cache_policy_lives_with_the_server(tmp_path): assert "cache-control" not in client.get("/assets/gone.js").headers -def test_request_that_never_touches_the_db_opens_no_session(monkeypatch): - opened = [] - original = db_session.registry.createfunc - - def counting(): - opened.append(1) - return original() - - monkeypatch.setattr(db_session.registry, "createfunc", counting) - # The harness binds an ambient session; drop it so what the registry holds here - # is only what a request opened. - db_session.remove() +def test_request_that_never_touches_the_db_opens_no_connection(monkeypatch): + from sqlalchemy.orm import Session + + connected = [] + original = Session._connection_for_bind + + def counting(self, *args, **kwargs): + connected.append(1) + return original(self, *args, **kwargs) + + monkeypatch.setattr(Session, "_connection_for_bind", counting) + # Outside a loop the registry scopes to None; an earlier test's leftover + # binding under that key would masquerade as ours below. + db_session.registry.clear() app = FastAPI(dependencies=[Depends(_release_db_session)]) @app.get("/plain") async def plain() -> dict[str, str]: return {} - @app.get("/touch") - async def touch() -> dict[str, str]: - db_session() - return {} - client = TestClient(app) - # The SPA/asset case: the boundary must not open a session to commit nothing. + # The SPA/asset case: the boundary binds a session for every request, but + # the session is lazy — serving a request that never touches the DB checks + # out no connection, and the commit is a no-op. The binding is gone after. assert client.get("/plain").status_code == 200 - assert not opened - # The API case: a session the request opened still gets committed + released. - assert client.get("/touch").status_code == 200 - assert opened == [1] + assert not connected assert not db_session.registry.has() diff --git a/backend/tests/test_usage_poll.py b/backend/tests/test_usage_poll.py index 9889abd0..527e88f0 100644 --- a/backend/tests/test_usage_poll.py +++ b/backend/tests/test_usage_poll.py @@ -53,19 +53,19 @@ async def fetch_usage(cls, connection, *, now=None): return _Fake -def _connection(email: str = "op@example.com"): +async def _connection(email: str = "op@example.com"): # poll_usage reads only account_id off the connection; the account row # must be real (the scrape carries its FK). from types import SimpleNamespace from druks.accounts.models import Account - return SimpleNamespace(account_id=Account.get_or_create(email).id) + return SimpleNamespace(account_id=(await Account.get_or_create(email)).id) async def _poll(*harnesses) -> list[dict[str, object]]: # poll_usage is the unit under test: fetch -> parse -> persist a UsageScrape. - connection = _connection() + connection = await _connection() return [await h.poll_usage(connection) for h in harnesses] @@ -77,19 +77,19 @@ async def test_successful_fetch_persists_per_harness(druks_db) -> None: lambda: _usage(plan_tier="prolite", five=_metric(61), weeks=(_metric(61),)), ), ) - druks_db.flush() + await druks_db.flush() assert [r["status"] for r in results] == ["recorded", "recorded"] assert all(r["parse_ok"] for r in results) - claude_row = UsageScrape.latest_for("claude", _connection().account_id) + claude_row = await UsageScrape.latest_for("claude", (await _connection()).account_id) assert claude_row is not None assert claude_row.five_hour_percent_left == 84 assert claude_row.weeks == [ {"percent_left": 52, "resets_at": None, "model": None}, ] - codex_row = UsageScrape.latest_for("codex", _connection().account_id) + codex_row = await UsageScrape.latest_for("codex", (await _connection()).account_id) assert codex_row is not None assert codex_row.plan_tier == "prolite" assert codex_row.weeks[0]["percent_left"] == 61 @@ -114,9 +114,9 @@ async def test_claude_weekly_windows_survive_parse_and_poll_in_order(druks_db) - ) await _poll(_harness("claude", lambda: parsed)) - druks_db.flush() + await druks_db.flush() - row = UsageScrape.latest_for("claude", _connection().account_id) + row = await UsageScrape.latest_for("claude", (await _connection()).account_id) assert row is not None assert [(week["percent_left"], week["model"]) for week in row.weeks] == [ (70, None), @@ -129,11 +129,11 @@ async def test_credential_error_records_error_snapshot(druks_db) -> None: _harness("claude", lambda: _usage(ok=False, error="token_expired")), _harness("codex", lambda: _usage(ok=False, error="no_credentials")), ) - druks_db.flush() + await druks_db.flush() assert all(r["status"] == "recorded" for r in results) assert all(not r["parse_ok"] for r in results) - claude_row = UsageScrape.latest_for("claude", _connection().account_id) + claude_row = await UsageScrape.latest_for("claude", (await _connection()).account_id) assert claude_row is not None assert claude_row.parse_ok is False assert claude_row.error == "token_expired" @@ -145,10 +145,10 @@ def boom() -> ParsedUsage: raise RuntimeError("boom") results = await _poll(_harness("claude", boom), _harness("codex", boom)) - druks_db.flush() + await druks_db.flush() assert all(r["status"] == "errored" and r["error"] == "crashed" for r in results) - row = UsageScrape.latest_for("claude", _connection().account_id) + row = await UsageScrape.latest_for("claude", (await _connection()).account_id) assert row is not None assert row.parse_ok is False @@ -165,9 +165,9 @@ async def test_snapshot_persists_unlimited_flag(druks_db) -> None: ), ) ) - druks_db.flush() + await druks_db.flush() - row = UsageScrape.latest_for("codex", _connection().account_id) + row = await UsageScrape.latest_for("codex", (await _connection()).account_id) assert row is not None assert row.unlimited is True @@ -175,11 +175,11 @@ async def test_snapshot_persists_unlimited_flag(druks_db) -> None: async def test_two_accounts_of_one_harness_snapshot_independently(druks_db) -> None: snapshots = iter([_usage(five=_metric(84)), _usage(five=_metric(30))]) fake = _harness("claude", lambda: next(snapshots)) - first, second = _connection("a@example.com"), _connection("b@example.com") + first, second = await _connection("a@example.com"), await _connection("b@example.com") await fake.poll_usage(first) await fake.poll_usage(second) - druks_db.flush() + await druks_db.flush() - assert UsageScrape.latest_for("claude", first.account_id).five_hour_percent_left == 84 - assert UsageScrape.latest_for("claude", second.account_id).five_hour_percent_left == 30 + assert (await UsageScrape.latest_for("claude", first.account_id)).five_hour_percent_left == 84 + assert (await UsageScrape.latest_for("claude", second.account_id)).five_hour_percent_left == 30 diff --git a/backend/tests/test_user_settings.py b/backend/tests/test_user_settings.py index d7434b12..09ec19d4 100644 --- a/backend/tests/test_user_settings.py +++ b/backend/tests/test_user_settings.py @@ -21,29 +21,30 @@ def session(druks_db): return druks_db -def test_get_lazy_creates_row_with_default_timezone(session): - row = UserSettings.get() - session.commit() +async def test_get_lazy_creates_row_with_default_timezone(session): + row = await UserSettings.get() + await session.commit() assert row.timezone == "UTC" -def test_harnesses_seeded_with_shipped_defaults(session): +async def test_harnesses_seeded_with_shipped_defaults(session): # init_db seeds one HarnessSettings row per registered harness. - claude = HarnessSettings.require("claude") + claude = await HarnessSettings.require("claude") assert (claude.model, claude.fast_mode, claude.effort, claude.timeout) == ( "claude-opus-4-7", False, "high", 1800, ) - assert HarnessSettings.require("codex").model == "gpt-5.5" - assert {harness.name for harness in HarnessSettings.all()} == {"claude", "codex"} + assert (await HarnessSettings.require("codex")).model == "gpt-5.5" + assert {harness.name for harness in await HarnessSettings.all()} == {"claude", "codex"} -def test_harness_update_persists(session): - HarnessSettings.require("claude").update(model="claude-sonnet-4-6", fast_mode=True) - session.commit() - claude = HarnessSettings.require("claude") +async def test_harness_update_persists(session): + claude = await HarnessSettings.require("claude") + await claude.update(model="claude-sonnet-4-6", fast_mode=True) + await session.commit() + claude = await HarnessSettings.require("claude") assert claude.model == "claude-sonnet-4-6" assert claude.fast_mode is True diff --git a/backend/tests/test_webhooks_framework.py b/backend/tests/test_webhooks_framework.py index 92a618a5..db114258 100644 --- a/backend/tests/test_webhooks_framework.py +++ b/backend/tests/test_webhooks_framework.py @@ -45,7 +45,7 @@ class Hook(Webhook): provider = "acme" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -70,7 +70,7 @@ class Hook(Mid): provider = "acme" category = "alerts" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -84,7 +84,7 @@ def test_concrete_without_provider_and_category_raises(): with pytest.raises(TypeError, match="provider"): class Hook(Webhook): - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -97,7 +97,7 @@ class Hook(Webhook): provider = "ignored" category = "ignored" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -114,7 +114,7 @@ class Hook(Webhook): provider = "alpha" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -131,7 +131,7 @@ class Hook(Webhook): provider = "tenant" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -155,7 +155,7 @@ class Hook(Webhook): provider = "late" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -187,7 +187,7 @@ class Hook(Webhook): provider = "dispatch" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -209,7 +209,7 @@ class Hook(Webhook): provider = "fallback" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -231,7 +231,7 @@ class Hook(Webhook): provider = "auth" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return False def get_action(self): @@ -247,7 +247,7 @@ class Hook(Webhook): provider = "rich" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): raise HTTPException(status_code=403, detail="Forbidden by policy.") def get_action(self): @@ -269,7 +269,7 @@ class Hook(Webhook): provider = "loose" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -293,7 +293,7 @@ class Hook(Webhook): provider = "dedup" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def get_action(self): @@ -327,7 +327,7 @@ class Hook(Webhook): provider = "flaky" category = "events" - def request_is_authentic(self): + async def request_is_authentic(self): return True def delivery_key(self): diff --git a/backend/tests/test_webhooks_slack.py b/backend/tests/test_webhooks_slack.py index 1371e264..401c15da 100644 --- a/backend/tests/test_webhooks_slack.py +++ b/backend/tests/test_webhooks_slack.py @@ -27,7 +27,7 @@ } -def _parked_notification(druks_db): +async def _parked_notification(druks_db): run = Run( id=str(uuid7()), kind="notifications.test", @@ -36,12 +36,12 @@ def _parked_notification(druks_db): input_requested_at=Base.utc_now(), ) druks_db.add(run) - druks_db.flush() - seed_dbos_status(druks_db, run.id, "parked") - destination = Destination.create( + await druks_db.flush() + await seed_dbos_status(druks_db, run.id, "parked") + destination = await Destination.create( name=f"slack-{run.id[-8:]}", kind="slack_webhook", url=_WEBHOOK_URL ) - notification = Notification.create( + notification = await Notification.create( destination_id=destination.id, reason="gate.parked", body="review the plan", @@ -117,7 +117,7 @@ def test_signature_verifier_against_a_known_answer(monkeypatch): async def test_unsigned_or_stale_requests_401_and_never_resume(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) body = _interactivity_body(encode_button(notification.correlation_token, "approve")) with _client(tmp_path) as client: no_headers = client.post("/_external/slack/interactivity/", content=body) @@ -138,14 +138,14 @@ async def test_unsigned_or_stale_requests_401_and_never_resume(tmp_path, druks_d assert replayed.status_code == 401 assert resume_spy == [] - assert Notification.get(notification.id).state == "pending" + assert (await Notification.get(notification.id)).state == "pending" for response in (no_headers, wrong_secret, replayed): assert _SIGNING_SECRET not in response.text assert notification.correlation_token not in response.text async def test_signed_click_routes_through_respond(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) body = _interactivity_body(encode_button(notification.correlation_token, "approve")) with _client(tmp_path) as client: response = client.post( @@ -157,15 +157,15 @@ async def test_signed_click_routes_through_respond(tmp_path, druks_db, resume_sp assert resume_spy == [{"id": run.id, "action": "approve", "answers": {}, "note": ""}] # The request's own session committed the transition; drop the ambient # (task-scoped) session's cached instance to read it. - ambient_db_session().expire_all() - assert Notification.get(notification.id).state == "acknowledged" + ambient_db_session().expunge_all() + assert (await Notification.get(notification.id)).state == "acknowledged" assert notification.correlation_token not in response.text assert _SIGNING_SECRET not in response.text async def test_dead_round_click_is_acknowledged_without_resume(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db) - notification.mark_acknowledged() + run, notification = await _parked_notification(druks_db) + await notification.mark_acknowledged() with _client(tmp_path) as client: body = _interactivity_body(encode_button(notification.correlation_token, "approve")) acknowledged = client.post( @@ -186,7 +186,7 @@ async def test_dead_round_click_is_acknowledged_without_resume(tmp_path, druks_d async def test_malformed_payloads_400_never_500_never_resume(tmp_path, druks_db, resume_spy): - run, notification = _parked_notification(druks_db) + run, notification = await _parked_notification(druks_db) malformed = [ b"not-a-form", urlencode({"payload": "not json"}).encode(), @@ -204,11 +204,11 @@ async def test_malformed_payloads_400_never_500_never_resume(tmp_path, druks_db, assert response.status_code == 400 assert resume_spy == [] - assert Notification.get(notification.id).state == "pending" + assert (await Notification.get(notification.id)).state == "pending" async def test_unknown_interactivity_type_is_acknowledged_unhandled(tmp_path, druks_db, resume_spy): - _parked_notification(druks_db) + await _parked_notification(druks_db) body = urlencode({"payload": json.dumps({"type": "view_submission"})}).encode() with _client(tmp_path) as client: diff --git a/backend/tests/test_workflow_identity.py b/backend/tests/test_workflow_identity.py index 93a27e0b..23aa0779 100644 --- a/backend/tests/test_workflow_identity.py +++ b/backend/tests/test_workflow_identity.py @@ -11,6 +11,7 @@ from druks.events.models import Event from druks.workflows import Gate, Workflow, _log_run_event, step, task from druks_field_notes.models import Note +from sqlalchemy import select @pytest.fixture(autouse=True) @@ -158,18 +159,18 @@ async def ping(self) -> None: ... assert "alpha.pinger" in captured -def test_lifecycle_event_stamps_the_declaring_app(druks_db): +async def test_lifecycle_event_stamps_the_declaring_app(druks_db): # The event's app derives from the run's kind through the registry — # never an argument, never a stored copy on the run. register_workflow_package("alpha_pkg", "alpha") flow = _workflow("Beacon", "alpha_pkg.workflows") run = Run(id="wf-identity-1", kind=flow.kind) druks_db.add(run) - druks_db.flush() + await druks_db.flush() - payload = _log_run_event(run, RunState.FINISHED, {"type": "note", "id": 1}) + payload = await _log_run_event(run, RunState.FINISHED, {"type": "note", "id": 1}, None) - event = druks_db.query(Event).filter_by(type="workflow.finished").one() + event = (await druks_db.scalars(select(Event).filter_by(type="workflow.finished"))).one() assert payload["run"] == run.id assert event.app == "alpha" @@ -178,14 +179,14 @@ def test_display_label_reads_the_local_kind(): assert get_display_label("field_notes.summarize") == "Summarize" -def test_a_declared_subject_answers_off_the_workflow_and_off_a_run(): +async def test_a_declared_subject_answers_off_the_workflow_and_off_a_run(): # One word, two answers: the workflow says what kind of thing its runs are about, # a run of it says which one. register_workflow_package("alpha_pkg", "alpha") flow = _workflow("Sweep", "alpha_pkg.workflows", subject=Note) assert flow.subject is Note - assert flow().subject is None # a run with no subject has none to resolve + assert await flow().subject is None # a run with no subject has none to resolve def test_a_workflow_about_nothing_says_so_by_silence():