From 86dbd3fba01f51a63a942b849b8d95085994c1e2 Mon Sep 17 00:00:00 2001 From: Paulo Date: Wed, 9 Sep 2026 19:17:04 +0200 Subject: [PATCH 1/2] Record useful outputs through saved artifacts (DRU-506) --- backend/druks/agents.py | 30 ++++++- backend/druks/durable/models.py | 30 ++++++- backend/tests/test_agents.py | 43 ++++++++++ backend/tests/test_durable_sdk.py | 52 ++++++++++++ backend/tests/test_output_activity.py | 111 ++++++++++++++++++++++++++ 5 files changed, 262 insertions(+), 4 deletions(-) create mode 100644 backend/tests/test_output_activity.py diff --git a/backend/druks/agents.py b/backend/druks/agents.py index 5200978d..60f6c18c 100644 --- a/backend/druks/agents.py +++ b/backend/druks/agents.py @@ -102,6 +102,10 @@ def get_artifact(self) -> dict[str, str]: # it after the call. Empty unless the contract produces a reviewable document. return {} + def get_activity(self) -> dict[str, str]: + """Return a kind and optional summary for the saved result. Empty opts out.""" + return {} + @dataclass(frozen=True) class Agent: @@ -354,6 +358,26 @@ async def _run( f"agent {self.id!r} returned a payload that fails " f"{contract.__name__}: {error}" ) from error + activity = output.get_activity() + spec = output.get_artifact() + if not isinstance(activity, dict) or ( + activity + and ( + not isinstance(activity.get("kind"), str) + or not activity["kind"].strip() + or set(activity) - {"kind", "summary"} + or any(not isinstance(value, str) for value in activity.values()) + ) + ): + raise WorkflowError( + f"{contract.__name__}.get_activity() returned an invalid contract. " + "Return an empty dict or a kind with an optional string summary." + ) + if activity and not spec: + raise WorkflowError( + f"{contract.__name__} declares activity without an artifact. " + "Implement get_artifact() for the saved result." + ) if workspace_files: await runner.save_files( workspace_files, @@ -365,6 +389,8 @@ async def _run( await AgentCall.fail(engine, call_id=call_id, error=error) raise - if spec := output.get_artifact(): - await Artifact.record(call_dir=artifact_dir / call_id, call_id=call_id, **spec) + if spec: + await Artifact.record( + call_dir=artifact_dir / call_id, call_id=call_id, activity=activity, **spec + ) return output.to_result() diff --git a/backend/druks/durable/models.py b/backend/druks/durable/models.py index 78492398..f70f0d6d 100644 --- a/backend/druks/durable/models.py +++ b/backend/druks/durable/models.py @@ -11,6 +11,7 @@ from sqlalchemy.orm import Mapped, column_property, mapped_column, relationship, selectinload from druks.accounts.models import Account +from druks.apps.registry import workflows from druks.core.models import Uuid7Pk from druks.database import db_session, get_session from druks.durable.dbos_state import ( @@ -29,6 +30,7 @@ WorkflowEvent, ) from druks.durable.exceptions import AgentCallNotFound +from druks.events.models import Event from druks.harnesses.artifacts import normalize_token_usage from druks.models import Base from druks.notifications.models import Notification @@ -655,7 +657,14 @@ class Artifact(Base, Uuid7Pk): @classmethod async def record( - cls, *, call_dir: Path, call_id: str, kind: str, title: str, content: str + cls, + *, + call_dir: Path, + call_id: str, + kind: str, + title: str, + content: str, + activity: dict[str, str] | None = None, ) -> 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 @@ -664,11 +673,28 @@ async def record( call_dir.mkdir(parents=True, exist_ok=True) (call_dir / name).write_text(content) session = db_session() - await session.execute( + artifact_id = await session.scalar( pg_insert(cls) .values(agent_call_id=call_id, kind=kind, title=title, path=name) .on_conflict_do_nothing(index_elements=["agent_call_id"]) + .returning(cls.id) ) + if artifact_id and activity: + call = await AgentCall.get(call_id) + run = call.run + await Event.emit( + type=activity["kind"], + subject=await run.get_subject(), + label=run.subject_label, + app=workflows.get(run.kind).app, + payload={ + "run": run.id, + "kind": run.kind, + "agent_call_id": call.id, + "artifact_id": artifact_id, + **{key: value for key, value in activity.items() if key == "summary"}, + }, + ) await session.flush() @classmethod diff --git a/backend/tests/test_agents.py b/backend/tests/test_agents.py index 2db550fd..9234810e 100644 --- a/backend/tests/test_agents.py +++ b/backend/tests/test_agents.py @@ -9,6 +9,7 @@ from druks.accounts.models import Account from druks.database import db_session from druks.durable import AgentCall, WorkflowError +from druks.events.models import Event from druks.files import File from druks.sandbox.exceptions import SandboxDownloadError from druks.sandbox.models import SandboxIdentity, SecretRef @@ -1007,3 +1008,45 @@ async def fake_ephemeral(self, **_kwargs): assert result == DummyOutput(ok=True) assert resumed == ["host-crashed"] assert [row.id for row in await _identities("wf-9")] == [identity.id] + + +@pytest.mark.parametrize( + "activity", + [ + None, + [], + "review.completed", + {"summary": "Ready"}, + {"kind": ""}, + {"kind": "review.completed", "summary": 1}, + {"kind": "review.completed", "copy": "Ready"}, + ], +) +async def test_invalid_activity_contract_records_no_success( + druks_db, tmp_path, monkeypatch, current_run, activity +): + sandbox = _patch_runtime(monkeypatch, tmp_path, {"ok": True}) + _patch_ephemeral(monkeypatch, sandbox) + monkeypatch.setattr(DummyOutput, "get_activity", lambda self: activity) + + with pytest.raises(WorkflowError, match="invalid contract"): + await DUMMY_AGENT._run(workflow_id="wf-9") + + assert not list(await db_session().scalars(select(agents.Artifact))) + assert not list(await db_session().scalars(select(Event))) + calls = await AgentCall.list_for_run("wf-9") + assert len(calls) == 1 + assert calls[0].status == "failed" + + +async def test_activity_requires_an_artifact(druks_db, tmp_path, monkeypatch, current_run): + sandbox = _patch_runtime(monkeypatch, tmp_path, {"ok": True}) + _patch_ephemeral(monkeypatch, sandbox) + monkeypatch.setattr(DummyOutput, "get_activity", lambda self: {"kind": "review.completed"}) + + with pytest.raises(WorkflowError, match="without an artifact"): + await DUMMY_AGENT._run(workflow_id="wf-9") + + assert not list(await db_session().scalars(select(agents.Artifact))) + assert not list(await db_session().scalars(select(Event))) + assert (await AgentCall.list_for_run("wf-9"))[0].status == "failed" diff --git a/backend/tests/test_durable_sdk.py b/backend/tests/test_durable_sdk.py index 89921916..5835afc3 100644 --- a/backend/tests/test_durable_sdk.py +++ b/backend/tests/test_durable_sdk.py @@ -13,6 +13,7 @@ from druks.durable import FatalError, Run, RunState from druks.durable.dbos_state import workflow_status from druks.durable.engine import configure_engine, init_dbos, launch, shutdown +from druks.durable.models import Artifact from druks.events.models import Event from druks.models import StoredSubject from druks.signals import subscribe @@ -53,6 +54,16 @@ class RepoCfg(BaseModel): repo: str +class ReviewResult(AgentOutput): + action: str + + def get_artifact(self) -> dict[str, str]: + return {"kind": "markdown", "title": "Review", "content": self.action} + + def get_activity(self) -> dict[str, str]: + return {"kind": "review.completed", "summary": self.action} + + SINK: list[str] = [] TASK_RETRY_ATTEMPTS = 0 STEP_RETRY_ATTEMPTS = 0 @@ -1383,3 +1394,44 @@ async def run(self) -> None: assert {event.payload["failure"] for event in failures} == {"Source unavailable"} finally: workflows._items.pop(FailingAttempt.kind) + + +async def test_output_activity_survives_completed_step_replay(rt, monkeypatch): + calls = [] + held = [] + monkeypatch.setattr( + "druks.sandbox.client.Client.ephemeral", _fake_ephemeral_returning("reviewed", calls, held) + ) + monkeypatch.setattr("druks.agents.render_prompt", _fake_render) + completed = [] + + class OutputFlow(Workflow): + subject = Widget + + async def run_multistep(self) -> None: + for _ in range(2): + await rt.AgentFlow.DECIDER(contract=ReviewResult, body="x") + completed.append(self.workflow_id) + if len(completed) == 1: + raise asyncio.CancelledError + + workflow_id = await OutputFlow.start(subject=Widget(id=7)) + try: + await _wait_for(rt.engine, workflow_id, lambda run: len(completed) == 1) + await DBOS.resume_workflow_async(workflow_id) + await _wait_for(rt.engine, workflow_id, lambda run: run.state == RunState.FINISHED) + assert len(completed) == 2 + assert len(calls) == 2 + async with get_session(rt.engine) as session: + events = list(await session.scalars(select(Event).filter_by(type="review.completed"))) + artifacts = list(await session.scalars(select(Artifact))) + assert len(events) == len(artifacts) == 2 + assert {event.payload["artifact_id"] for event in events} == { + artifact.id for artifact in artifacts + } + assert {event.payload["agent_call_id"] for event in events} == { + call["call_id"] for call in calls + } + assert {event.payload["run"] for event in events} == {workflow_id} + finally: + workflows._items.pop(OutputFlow.kind) diff --git a/backend/tests/test_output_activity.py b/backend/tests/test_output_activity.py new file mode 100644 index 00000000..1975c976 --- /dev/null +++ b/backend/tests/test_output_activity.py @@ -0,0 +1,111 @@ +from unittest.mock import AsyncMock + +import pytest +from conftest import installation_key +from druks.database import db_session +from druks.durable.models import AgentCall, Artifact +from druks.events.models import Event +from druks.testing import seed_run +from druks_field_notes.models import Note +from druks_field_notes.workflows import Summarize +from sqlalchemy import select + + +@pytest.fixture +async def output_calls(druks_db): + db_session.registry.set(druks_db) + note = await Note.create(body="The reviewed work") + run = await seed_run(druks_db, kind=Summarize.kind, subject=note) + key = await installation_key() + calls = [ + AgentCall( + id=f"output-{number}", + run_id=run.id, + agent="field_notes.summarize", + model="test", + sandbox_host_id="test", + api_key_id=key.id, + ) + for number in (1, 2) + ] + druks_db.add_all(calls) + await druks_db.flush() + return note, run, calls + + +async def test_output_persistence_keeps_one_event_per_call(druks_db, output_calls, tmp_path): + db_session.registry.set(druks_db) + note, run, calls = output_calls + for call in calls: + for _ in range(2): + await Artifact.record( + call_dir=tmp_path / call.id, + call_id=call.id, + kind="markdown", + title="Review", + content="No unresolved findings.", + activity={"kind": "review.completed", "summary": "No unresolved findings."}, + ) + await druks_db.commit() + artifacts = list(await druks_db.scalars(select(Artifact).order_by(Artifact.agent_call_id))) + events = list(await druks_db.scalars(select(Event).order_by(Event.id))) + assert len(artifacts) == len(events) == 2 + assert {event.app for event in events} == {"field_notes"} + assert {event.subject_id for event in events} == {str(note.id)} + assert {event.subject_label for event in events} == {note.label} + for call, artifact, event in zip(calls, artifacts, events, strict=True): + assert event.type == "review.completed" + assert event.payload == { + "run": run.id, + "kind": run.kind, + "agent_call_id": call.id, + "artifact_id": artifact.id, + "summary": "No unresolved findings.", + } + + +async def test_artifact_without_activity_creates_no_event(druks_db, output_calls, tmp_path): + db_session.registry.set(druks_db) + _, _, calls = output_calls + await Artifact.record( + call_dir=tmp_path, + call_id=calls[0].id, + kind="markdown", + title="Internal result", + content="A working note.", + ) + assert await Artifact.get_for_call(calls[0].id) + assert not list(await druks_db.scalars(select(Event))) + + +async def test_event_failure_rolls_back_the_artifact(druks_db, output_calls, tmp_path, monkeypatch): + db_session.registry.set(druks_db) + _, _, calls = output_calls + call_id = calls[0].id + with monkeypatch.context() as patched: + patched.setattr(Event, "emit", AsyncMock(side_effect=RuntimeError("Event insert failed"))) + with pytest.raises(RuntimeError, match="Event insert failed"): + async with druks_db.begin_nested(): + await Artifact.record( + call_dir=tmp_path, + call_id=call_id, + kind="markdown", + title="Review", + content="Reviewed.", + activity={"kind": "review.completed"}, + ) + assert not await Artifact.get_for_call(call_id) + assert not list(await druks_db.scalars(select(Event))) + + await Artifact.record( + call_dir=tmp_path, + call_id=call_id, + kind="markdown", + title="Review", + content="Reviewed.", + activity={"kind": "review.completed"}, + ) + assert await Artifact.get_for_call(call_id) + events = list(await druks_db.scalars(select(Event))) + assert len(events) == 1 + assert "summary" not in events[0].payload From b6340f73a8adb5ad77baa09fe4a5302b56cbcdea Mon Sep 17 00:00:00 2001 From: Paulo Date: Wed, 9 Sep 2026 19:22:05 +0200 Subject: [PATCH 2/2] Load output fixtures from saved rows (DRU-506) --- backend/tests/test_durable_sdk.py | 3 ++- backend/tests/test_output_activity.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/tests/test_durable_sdk.py b/backend/tests/test_durable_sdk.py index 5835afc3..8ad0301d 100644 --- a/backend/tests/test_durable_sdk.py +++ b/backend/tests/test_durable_sdk.py @@ -1396,7 +1396,8 @@ async def run(self) -> None: workflows._items.pop(FailingAttempt.kind) -async def test_output_activity_survives_completed_step_replay(rt, monkeypatch): +async def test_output_activity_survives_completed_step_replay(rt, monkeypatch, tmp_path): + monkeypatch.setenv("DRUKS_DATA_DIR", str(tmp_path)) calls = [] held = [] monkeypatch.setattr( diff --git a/backend/tests/test_output_activity.py b/backend/tests/test_output_activity.py index 1975c976..b97d1c2f 100644 --- a/backend/tests/test_output_activity.py +++ b/backend/tests/test_output_activity.py @@ -30,6 +30,7 @@ async def output_calls(druks_db): ] druks_db.add_all(calls) await druks_db.flush() + druks_db.expunge_all() return note, run, calls