From 6efa671b1cb83735c83a33bcea7e3212f0b12aaa Mon Sep 17 00:00:00 2001 From: DSCmatter Date: Tue, 18 Aug 2026 11:54:03 +0530 Subject: [PATCH 1/4] feat: implement phase 2 causal event graph - add PostgreSQL snapshots table and graph indexes - implement explicit cross-agent causal parent assignment - add PostgreSQL-backed ancestors(event_id) traversal - add real three-agent Phase 2 integration coverage - validate planner and worker branches against the fixture - update CI, setup, contribution, and Neon verification docs --- .github/workflows/ci.yml | 2 +- CONTRIBUTING.md | 5 +- GETTING_STARTED.md | 33 ++- README.md | 30 +++ TEST.md | 541 +++++++++++++++++++++------------------ core/__init__.py | 5 + core/graph.py | 66 +++++ storage/postgres.py | 46 +++- tests/test_phase2.py | 422 ++++++++++++++++++++++++++++++ 9 files changed, 886 insertions(+), 264 deletions(-) create mode 100644 core/__init__.py create mode 100644 core/graph.py create mode 100644 tests/test_phase2.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 76066ec..1009155 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -73,4 +73,4 @@ jobs: run: uv sync --locked - name: Run PostgreSQL integration test - run: uv run pytest tests/test_postgres_integration.py -m integration -q + run: uv run pytest tests/test_postgres_integration.py tests/test_phase2.py -m integration -q diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f97ab10..4a502cc 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -117,9 +117,12 @@ git rebase --abort - For database changes, also run the Neon-backed integration test: ```powershell - uv run --env-file .env pytest tests/test_postgres_integration.py -m integration -q + uv run --env-file .env pytest tests/test_postgres_integration.py tests/test_phase2.py -m integration -q ``` + This covers both the Phase 1 PostgreSQL capture paths and the Phase 2 + schema, explicit cross-agent merge, and ancestor query. + ## CI gate Every push to `main` and every pull request runs the GitHub Actions workflow diff --git a/GETTING_STARTED.md b/GETTING_STARTED.md index 4950d78..4d2bda4 100644 --- a/GETTING_STARTED.md +++ b/GETTING_STARTED.md @@ -1,8 +1,8 @@ # Getting Started Agent-Casuality is a causal event-capture SDK for branching multi-agent -systems. Phase 1 records model calls, tool calls, memory operations, and -agent spawning in PostgreSQL. +systems. Phase 1 records execution events, and Phase 2 makes their causal +agent/event graph queryable in PostgreSQL. ## Prerequisites @@ -47,16 +47,17 @@ powershell -NoProfile -ExecutionPolicy Bypass -File scripts/check.ps1 The script runs pytest, Ruff, and ty. If `.env` exists, it loads the file so the PostgreSQL integration test runs as well. -Run the real PostgreSQL integration test with the `.env` file loaded: +Run the real Phase 1 and Phase 2 PostgreSQL integration tests with the `.env` +file loaded: ```powershell -uv run --env-file .env pytest tests/test_postgres_integration.py -m integration -q +uv run --env-file .env pytest tests/test_postgres_integration.py tests/test_phase2.py -m integration -q ``` Expected result: ```text -1 passed +3 passed ``` To run the complete suite with PostgreSQL enabled: @@ -147,20 +148,26 @@ transaction-scoped PostgreSQL advisory lock. ## Verify the data in Neon Open the Neon SQL Editor for the same branch used by `DATABASE_URL`. Run the -queries in [TEST.md](TEST.md). They verify: +queries in [TEST.md](TEST.md). They verify the Phase 1 capture data and the +Phase 2 PostgreSQL graph: -- each Phase 1 run has three agents and seven events +- required Phase 2 tables, columns, indexes, and foreign keys +- expected agent and event counts for the integration runs - workers reference the planner and their spawn events -- tool calls and results are linked +- worker model, tool-call, and tool-result branches are linked - the planner merge preserves both worker result IDs +- graph ancestors include both worker branches +- every causal parent resolves to a real event - no agent has duplicate logical sequence numbers -The event count query uses `COUNT(DISTINCT e.id)`. Without `DISTINCT`, the -join between agents and events can report 21 instead of the actual 7 events. +The event count query uses `COUNT(DISTINCT ...)` because joining agents and +events multiplies rows. `TEST.md` explains the purpose and expected result of +each query, including why logical sequence numbers must not be used as causal +edges. ## Current scope -Phase 1 is complete. It includes: +Phase 1 and Phase 2 are complete. They include: - `Event` and thread-safe `AgentClock` - Anthropic `messages.create` capture @@ -168,8 +175,10 @@ Phase 1 is complete. It includes: - captured memory `get`, `set`, and `delete` - agent spawning with `spawned_at_event_id` - in-memory and PostgreSQL event/agent stores +- explicit cross-agent causal-parent assignment +- PostgreSQL-backed `ancestors(event_id)` queries -Phase 2+ features such as graph queries, state reconstruction, snapshots, +Phase 3+ features such as state reconstruction, snapshot creation, provenance traversal, replay, and minimal slicing are intentionally not yet implemented. diff --git a/README.md b/README.md index 3277cd0..08646e0 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,36 @@ # Agent-Casuality Causal Debugging for Branching Multi-Agent Systems +## Current implementation + +Phase 1 captures model calls, tool calls/results, memory operations, and agent +spawns. Phase 2 stores the resulting agent/event graph in PostgreSQL, supports +explicit cross-agent merge parents, and queries event ancestors. + +The Phase 2 graph uses `causal_parent_ids` as the source of dependency edges. +`logical_seq` is only the per-agent logical ordering value; it is not inferred +as a causal relationship. State reconstruction, snapshots, slicing, replay, +and provenance belong to later phases and are not implemented yet. + +## Install and run + +```powershell +uv sync +.\scripts\check.ps1 +``` + +For the real PostgreSQL scenario, put `DATABASE_URL` in a local `.env` file +and use a dedicated Neon branch or test database: + +```powershell +uv run --env-file .env pytest tests/test_postgres_integration.py tests/test_phase2.py -m integration -q +``` + +The integration tests create the schema through the existing PostgreSQL store, +capture the planner/worker scenario, assign the explicit merge parents, and +query `ancestors()` against PostgreSQL. See [GETTING_STARTED.md](GETTING_STARTED.md) +for setup and [TEST.md](TEST.md) for Neon SQL Editor verification queries. + ## Day-zero fixture Use `uv` to run the fixture so everyone gets the same Python entrypoint: diff --git a/TEST.md b/TEST.md index d85f91a..1f2f745 100644 --- a/TEST.md +++ b/TEST.md @@ -1,135 +1,181 @@ -## Testing Phase 1 +# Verify Phase 1 and Phase 2 in PostgreSQL -The commands below verify the implementation locally. The database command -also runs the real PostgreSQL integration test against the database in -`.env`. +Run the checks first so the database contains a fresh real run: -```pwsh +```powershell uv sync .\scripts\check.ps1 ``` -Expected result: tests, Ruff, and ty all pass. If `.env` contains a valid -`DATABASE_URL`, the PostgreSQL integration test runs automatically; otherwise -it is skipped. +The PostgreSQL tests use the `DATABASE_URL` from `.env`. They create the +schema if needed and leave test rows behind, so use a dedicated Neon branch +or database. -### Check the PostgreSQL console on Neon +The queries below are intended for the Neon SQL Editor. They do not modify +data. Because each test run uses generated UUIDs, the queries select the +latest run by `started_at` instead of hard-coding IDs. -#### 1. Confirm captured runs, agents, and event counts +## 1. Confirm the Phase 2 tables exist -Purpose: confirms that the integration test created runs and that each run -contains the expected three agents and seven Phase 1 events. `DISTINCT` is -important because joining agents and events otherwise multiplies the count. +Purpose: verifies that the existing Phase 1 schema and the Phase 2 +`snapshots` table were created. + +```sql +SELECT table_name +FROM information_schema.tables +WHERE table_schema = 'public' + AND table_name IN ('runs', 'agents', 'events', 'snapshots') +ORDER BY table_name; +``` + +Expected result: four rows: `agents`, `events`, `runs`, and `snapshots`. + +## 2. Inspect the required column types + +Purpose: confirms that IDs use UUIDs, causal parents remain a PostgreSQL UUID +array, event payloads use JSONB, and snapshot state uses JSONB. + +```sql +SELECT table_name, column_name, udt_name +FROM information_schema.columns +WHERE table_schema = 'public' + AND table_name IN ('runs', 'agents', 'events', 'snapshots') +ORDER BY table_name, ordinal_position; +``` + +Expected result: the required Phase 2 columns are present. In particular: + +- `events.id`, `events.run_id`, `events.agent_id` are `uuid`; +- `events.causal_parent_ids` is `_uuid` (PostgreSQL’s `uuid[]` type); +- `events.payload` and `snapshots.state` are `jsonb`; +- `snapshots.state_hash` is `text`. + +## 3. Confirm indexes and relationships + +Purpose: verifies the indexes used for event ordering, run queries, snapshots, +and retry idempotency. + +```sql +SELECT indexname, indexdef +FROM pg_indexes +WHERE schemaname = 'public' + AND indexname IN ( + 'idx_events_agent_seq', + 'idx_events_run_seq', + 'idx_snapshots_agent_seq', + 'idx_events_idempotency' + ) +ORDER BY indexname; +``` + +Expected result: four rows with those index names. + +Purpose: verifies the foreign keys, including +`agents.spawned_at_event_id -> events.id`. + +```sql +SELECT + kcu.table_name, + kcu.column_name, + ccu.table_name AS referenced_table, + ccu.column_name AS referenced_column +FROM information_schema.table_constraints tc +JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema +JOIN information_schema.constraint_column_usage ccu + ON tc.constraint_name = ccu.constraint_name + AND tc.table_schema = ccu.table_schema +WHERE tc.table_schema = 'public' + AND tc.constraint_type = 'FOREIGN KEY' +ORDER BY kcu.table_name, kcu.column_name; +``` + +Expected result includes these relationships: + +- `agents.run_id -> runs.id`; +- `agents.parent_agent_id -> agents.id`; +- `agents.spawned_at_event_id -> events.id`; +- `events.run_id -> runs.id`; +- `events.agent_id -> agents.id`; +- `snapshots.run_id -> runs.id`; +- `snapshots.agent_id -> agents.id`. + +## 4. Find the latest test runs + +Purpose: identifies the rows created by the Phase 1 and Phase 2 integration +tests without assuming fixed UUIDs. ```sql SELECT r.id AS run_id, r.name, + r.started_at, COUNT(DISTINCT a.id) AS agents, COUNT(DISTINCT e.id) AS events FROM runs r LEFT JOIN agents a ON a.run_id = r.id LEFT JOIN events e ON e.run_id = r.id -WHERE r.name = 'phase-1' -GROUP BY r.id, r.name -ORDER BY MAX(r.started_at) DESC; +WHERE r.name IN ('phase-1', 'phase-2', 'sequence') +GROUP BY r.id, r.name, r.started_at +ORDER BY r.started_at DESC; ``` -Expected result: one row per test run, with `agents = 3` and `events = 7`. -If the test was run twice, two `phase-1` rows are expected. - -RESULT: -```json -[{ - "run_id": "076e3473-fc4e-4c36-9441-698aa6b258e3", - "name": "phase-1", - "agents": 3, - "events": 7 -}, { - "run_id": "12e5846b-26fc-4de1-b054-76a4194fd0eb", - "name": "phase-1", - "agents": 3, - "events": 7 -}] -``` +Expected result for a fresh complete run: + +- `phase-1`: 3 agents and 7 events; +- `phase-2`: 3 agents and at least 9 events; +- `sequence`: 1 agent and 2 events. -#### 2. Confirm the planner-to-worker spawn relationships +The Phase 2 test also creates a shared-ancestor branch, so its event count is +higher than the minimum nine. -Purpose: verifies that the planner is the parent of both workers and that -each child stores the ID of its parent `agent_spawn` event in -`spawned_at_event_id`. +## 5. Confirm planner-to-worker relationships + +Purpose: verifies that the planner owns both workers and each worker stores +the exact event that spawned it in `spawned_at_event_id`. ```sql +WITH latest_phase2 AS ( + SELECT id + FROM runs + WHERE name = 'phase-2' + ORDER BY started_at DESC + LIMIT 1 +) SELECT a.id, a.role, a.parent_agent_id, a.spawned_at_event_id, + spawn.event_type AS spawned_event_type, a.lamport_offset FROM agents a -JOIN runs r ON r.id = a.run_id -WHERE r.name = 'phase-1' -ORDER BY a.created_at; +JOIN latest_phase2 r ON r.id = a.run_id +LEFT JOIN events spawn ON spawn.id = a.spawned_at_event_id +ORDER BY a.role, a.created_at; ``` -Expected result per run: - -- one `planner` with null `parent_agent_id` and `spawned_at_event_id` -- one `worker-1` and one `worker-2` -- both workers have the planner ID as `parent_agent_id` -- both workers have non-null `spawned_at_event_id` -- worker Lamport offsets match their spawn sequence numbers - -RESULT: -```json -[{ - "id": "d9a13702-5452-4a19-819e-8530400760cc", - "role": "planner", - "parent_agent_id": null, - "spawned_at_event_id": null, - "lamport_offset": 0 -}, { - "id": "18e7a9a3-5c2a-4ca2-b08b-a5bdd94b5752", - "role": "worker-1", - "parent_agent_id": "d9a13702-5452-4a19-819e-8530400760cc", - "spawned_at_event_id": "11de9524-9b4c-46c7-998b-b8080b163498", - "lamport_offset": 1 -}, { - "id": "1c62667d-5555-4c55-b655-eb6196fa70fb", - "role": "worker-2", - "parent_agent_id": "d9a13702-5452-4a19-819e-8530400760cc", - "spawned_at_event_id": "58d03d69-90e8-4d69-89f2-f363864d6725", - "lamport_offset": 2 -}, { - "id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "role": "planner", - "parent_agent_id": null, - "spawned_at_event_id": null, - "lamport_offset": 0 -}, { - "id": "28571a75-678d-4e08-8853-5288d731f30c", - "role": "worker-1", - "parent_agent_id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "spawned_at_event_id": "a69fbe1c-91dc-47c8-8cee-a42e691acf49", - "lamport_offset": 1 -}, { - "id": "0a212b0a-dd31-447c-a58a-5ed063471fba", - "role": "worker-2", - "parent_agent_id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "spawned_at_event_id": "b4b6b4aa-653b-4e9a-9065-34d9fdf09efc", - "lamport_offset": 2 -}] -``` +Expected result: one `planner`, one `researcher`, and one `coder`. +The planner has null parent/spawn fields. Both workers reference the planner, +have non-null `spawned_at_event_id`, and their referenced event type is +`agent_spawn`. -#### 3. Inspect event types, logical sequences, and causal parents +## 6. Inspect the real event branches -Purpose: verifies that all expected events were persisted and that causal -relationships are represented by `causal_parent_ids`, not by wall-clock -ordering. +Purpose: confirms that the worker branches contain model, tool-call, and +tool-result events, and that the result points to its tool call through +`causal_parent_ids`. ```sql +WITH latest_phase2 AS ( + SELECT id + FROM runs + WHERE name = 'phase-2' + ORDER BY started_at DESC + LIMIT 1 +) SELECT - e.agent_id, a.role, e.id, e.logical_seq, @@ -138,171 +184,158 @@ SELECT e.payload FROM events e JOIN agents a ON a.id = e.agent_id -JOIN runs r ON r.id = e.run_id -WHERE r.name = 'phase-1' -ORDER BY a.role, e.logical_seq; +JOIN latest_phase2 r ON r.id = e.run_id +ORDER BY a.role, e.logical_seq, e.id; ``` -Expected result per run: seven events—two `agent_spawn` events, two worker -`tool_call` events, two linked `tool_result` events, and one planner -`model_call` merge event. Each tool result should contain its tool call ID in -`causal_parent_ids`. +Expected result: the `researcher` and `coder` branches each contain +`model_call -> tool_call -> tool_result`. The worker model calls reference +their spawn events, tool calls reference the model calls, and tool results +reference the tool calls. The planner contains two spawn events and a merge +event. -RESULT: -```json -[{ - "agent_id": "d9a13702-5452-4a19-819e-8530400760cc", - "role": "planner", - "id": "11de9524-9b4c-46c7-998b-b8080b163498", - "logical_seq": 1, - "event_type": "agent_spawn", - "causal_parent_ids": "{}", - "payload": "{\"role\": \"worker-1\", \"child_agent_id\": \"18e7a9a3-5c2a-4ca2-b08b-a5bdd94b5752\"}" -}, { - "agent_id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "role": "planner", - "id": "a69fbe1c-91dc-47c8-8cee-a42e691acf49", - "logical_seq": 1, - "event_type": "agent_spawn", - "causal_parent_ids": "{}", - "payload": "{\"role\": \"worker-1\", \"child_agent_id\": \"28571a75-678d-4e08-8853-5288d731f30c\"}" -}, { - "agent_id": "d9a13702-5452-4a19-819e-8530400760cc", - "role": "planner", - "id": "58d03d69-90e8-4d69-89f2-f363864d6725", - "logical_seq": 2, - "event_type": "agent_spawn", - "causal_parent_ids": "{}", - "payload": "{\"role\": \"worker-2\", \"child_agent_id\": \"1c62667d-5555-4c55-b655-eb6196fa70fb\"}" -}, { - "agent_id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "role": "planner", - "id": "b4b6b4aa-653b-4e9a-9065-34d9fdf09efc", - "logical_seq": 2, - "event_type": "agent_spawn", - "causal_parent_ids": "{}", - "payload": "{\"role\": \"worker-2\", \"child_agent_id\": \"0a212b0a-dd31-447c-a58a-5ed063471fba\"}" -}, { - "agent_id": "d9a13702-5452-4a19-819e-8530400760cc", - "role": "planner", - "id": "1c9f4bd2-4d4e-4ebf-9349-dfddea6b9449", - "logical_seq": 3, - "event_type": "model_call", - "causal_parent_ids": "{21a24dfa-4e6c-4762-b3f4-d302568e9023,323c4acb-b1ad-4852-9230-fd7968b3e4bf}", - "payload": "{\"input\": [{\"role\": \"user\", \"content\": \"merge {'one': 'done'} {'two': 'done'}\"}], \"model\": \"test-model\", \"output\": {\"content\": \"merged\"}, \"latency_ms\": 0}" -}, { - "agent_id": "dda7e35a-6d0c-4eec-901b-52412e8b92a1", - "role": "planner", - "id": "6c823591-b622-43a2-90fa-751f97731f0e", - "logical_seq": 3, - "event_type": "model_call", - "causal_parent_ids": "{1dcae5aa-f478-4586-a0a8-9617a09aabc2,2ca4ee78-1658-4450-9fd5-59d71333893f}", - "payload": "{\"input\": [{\"role\": \"user\", \"content\": \"merge {'one': 'done'} {'two': 'done'}\"}], \"model\": \"test-model\", \"output\": {\"content\": \"merged\"}, \"latency_ms\": 0}" -}, { - "agent_id": "18e7a9a3-5c2a-4ca2-b08b-a5bdd94b5752", - "role": "worker-1", - "id": "14f289d0-ee44-40ce-86f2-8b9c99dc072c", - "logical_seq": 2, - "event_type": "tool_call", - "causal_parent_ids": "{}", - "payload": "{\"args\": [\"one\"], \"name\": \"inspect\", \"kwargs\": {}, \"invocation_id\": \"worker-one-tool\"}" -}, { - "agent_id": "28571a75-678d-4e08-8853-5288d731f30c", - "role": "worker-1", - "id": "143a75a3-e121-4c64-9036-b3b8dd5f49ed", - "logical_seq": 2, - "event_type": "tool_call", - "causal_parent_ids": "{}", - "payload": "{\"args\": [\"one\"], \"name\": \"inspect\", \"kwargs\": {}, \"invocation_id\": \"worker-one-tool\"}" -}, { - "agent_id": "18e7a9a3-5c2a-4ca2-b08b-a5bdd94b5752", - "role": "worker-1", - "id": "21a24dfa-4e6c-4762-b3f4-d302568e9023", - "logical_seq": 3, - "event_type": "tool_result", - "causal_parent_ids": "{14f289d0-ee44-40ce-86f2-8b9c99dc072c}", - "payload": "{\"output\": {\"one\": \"done\"}, \"invocation_id\": \"worker-one-tool\"}" -}, { - "agent_id": "28571a75-678d-4e08-8853-5288d731f30c", - "role": "worker-1", - "id": "1dcae5aa-f478-4586-a0a8-9617a09aabc2", - "logical_seq": 3, - "event_type": "tool_result", - "causal_parent_ids": "{143a75a3-e121-4c64-9036-b3b8dd5f49ed}", - "payload": "{\"output\": {\"one\": \"done\"}, \"invocation_id\": \"worker-one-tool\"}" -}, { - "agent_id": "1c62667d-5555-4c55-b655-eb6196fa70fb", - "role": "worker-2", - "id": "9d476664-2583-44b6-8e24-3fb652c5ee25", - "logical_seq": 3, - "event_type": "tool_call", - "causal_parent_ids": "{}", - "payload": "{\"args\": [\"two\"], \"name\": \"inspect\", \"kwargs\": {}, \"invocation_id\": \"worker-two-tool\"}" -}, { - "agent_id": "0a212b0a-dd31-447c-a58a-5ed063471fba", - "role": "worker-2", - "id": "bb0ee288-e727-4301-aba0-f91ba7db864d", - "logical_seq": 3, - "event_type": "tool_call", - "causal_parent_ids": "{}", - "payload": "{\"args\": [\"two\"], \"name\": \"inspect\", \"kwargs\": {}, \"invocation_id\": \"worker-two-tool\"}" -}, { - "agent_id": "0a212b0a-dd31-447c-a58a-5ed063471fba", - "role": "worker-2", - "id": "2ca4ee78-1658-4450-9fd5-59d71333893f", - "logical_seq": 4, - "event_type": "tool_result", - "causal_parent_ids": "{bb0ee288-e727-4301-aba0-f91ba7db864d}", - "payload": "{\"output\": {\"two\": \"done\"}, \"invocation_id\": \"worker-two-tool\"}" -}, { - "agent_id": "1c62667d-5555-4c55-b655-eb6196fa70fb", - "role": "worker-2", - "id": "323c4acb-b1ad-4852-9230-fd7968b3e4bf", - "logical_seq": 4, - "event_type": "tool_result", - "causal_parent_ids": "{9d476664-2583-44b6-8e24-3fb652c5ee25}", - "payload": "{\"output\": {\"two\": \"done\"}, \"invocation_id\": \"worker-two-tool\"}" -}] +## 7. Confirm the planner merge has both worker results as parents + +Purpose: verifies explicit cross-agent causal assignment and multiple parent +support. This checks actual parent rows rather than just counting array items. + +```sql +WITH latest_phase2 AS ( + SELECT id + FROM runs + WHERE name = 'phase-2' + ORDER BY started_at DESC + LIMIT 1 +), merge_event AS ( + SELECT e.* + FROM events e + JOIN latest_phase2 r ON r.id = e.run_id + JOIN agents a ON a.id = e.agent_id + WHERE a.role = 'planner' + AND e.event_type = 'model_call' + ORDER BY e.logical_seq DESC + LIMIT 1 +) +SELECT + m.id AS merge_event_id, + m.logical_seq AS merge_logical_seq, + p.id AS parent_event_id, + pa.role AS parent_role, + p.event_type AS parent_event_type, + p.logical_seq AS parent_logical_seq +FROM merge_event m +CROSS JOIN LATERAL unnest(m.causal_parent_ids) AS parents(parent_id) +JOIN events p ON p.id = parents.parent_id +JOIN agents pa ON pa.id = p.agent_id +ORDER BY pa.role; ``` -#### 4. Confirm the planner merge has both worker results as parents +Expected result: two rows for one merge event. The parent roles are +`researcher` and `coder`, and both parent event types are `tool_result`. +The merge’s `causal_parent_ids` are the real worker result IDs. -Purpose: verifies the multi-agent causal merge. The planner's final model -event must preserve both worker result IDs. +## 8. Query all ancestors of the merge event + +Purpose: runs the production PostgreSQL recursive query. It must include the +merge event itself and every event reachable through its causal parents. ```sql +WITH RECURSIVE latest_phase2 AS ( + SELECT id + FROM runs + WHERE name = 'phase-2' + ORDER BY started_at DESC + LIMIT 1 +), target AS ( + SELECT e.id, e.agent_id, e.logical_seq, e.causal_parent_ids + FROM events e + JOIN latest_phase2 r ON r.id = e.run_id + JOIN agents a ON a.id = e.agent_id + WHERE a.role = 'planner' + AND e.event_type = 'model_call' + ORDER BY e.logical_seq DESC + LIMIT 1 +), ancestors AS ( + SELECT id, agent_id, logical_seq, causal_parent_ids + FROM target + + UNION + + SELECT e.id, e.agent_id, e.logical_seq, e.causal_parent_ids + FROM events e + JOIN ancestors a ON e.id = ANY(a.causal_parent_ids) +) SELECT - e.id AS merge_event_id, - e.logical_seq, - e.causal_parent_ids, - e.payload->>'model' AS model -FROM events e -JOIN runs r ON r.id = e.run_id -WHERE r.name = 'phase-1' - AND e.event_type = 'model_call' -ORDER BY r.started_at DESC -LIMIT 1; + a.id, + agents.role, + a.logical_seq, + a.causal_parent_ids +FROM ancestors a +JOIN agents ON agents.id = a.agent_id +ORDER BY a.logical_seq, a.id; ``` -Expected result: one `model_call` with `model = 'test-model'`, and exactly two -IDs in `causal_parent_ids`. Those IDs should be the two worker -`tool_result` event IDs. - -RESULT: -```json -[{ - "merge_event_id": "6c823591-b622-43a2-90fa-751f97731f0e", - "logical_seq": 3, - "causal_parent_ids": "{1dcae5aa-f478-4586-a0a8-9617a09aabc2,2ca4ee78-1658-4450-9fd5-59d71333893f}", - "model": "test-model" -}] +Expected result: the merge event, both worker result branches, and the spawn +events reached through the worker model-call parents. No unrelated `sequence` +run or other agent appears. + +## 9. Confirm shared ancestors are deduplicated + +Purpose: verifies that two reachable paths to the same event return that event +once, not once per path. + +```sql +WITH RECURSIVE latest_phase2 AS ( + SELECT id + FROM runs + WHERE name = 'phase-2' + ORDER BY started_at DESC + LIMIT 1 +), target AS ( + SELECT e.id, e.causal_parent_ids + FROM events e + JOIN latest_phase2 r ON r.id = e.run_id + WHERE e.payload->>'shared_merge' = 'true' + LIMIT 1 +), ancestors AS ( + SELECT id, causal_parent_ids FROM target + UNION + SELECT e.id, e.causal_parent_ids + FROM events e + JOIN ancestors a ON e.id = ANY(a.causal_parent_ids) +) +SELECT + COUNT(*) AS ancestor_count, + COUNT(DISTINCT id) AS distinct_ancestor_count +FROM ancestors; ``` -#### 5. Confirm no duplicate logical sequences within an agent +Expected result: one row with `ancestor_count = 4` and +`distinct_ancestor_count = 4` for the shared-ancestor branch +(`shared_merge`, `left`, `right`, and the shared `root`). The equal counts +confirm that `UNION` prevented the shared root from appearing twice. -Purpose: checks the critical sequence invariant. Logical sequence numbers -may be equal across different agents, but must not collide for the same -agent. +## 10. Confirm every causal parent exists + +Purpose: checks that every stored causal-parent UUID resolves to an event. + +```sql +SELECT + child.id AS child_event_id, + parent_id, + child.causal_parent_ids +FROM events child +CROSS JOIN LATERAL unnest(child.causal_parent_ids) AS parents(parent_id) +LEFT JOIN events parent ON parent.id = parents.parent_id +WHERE parent.id IS NULL; +``` + +Expected result: no rows. Any row indicates a dangling causal dependency. + +## 11. Confirm no duplicate logical sequences within an agent + +Purpose: checks the critical ordering invariant. Logical sequences may match +across different agents, but cannot collide within one agent. ```sql SELECT agent_id, logical_seq, COUNT(*) @@ -311,6 +344,20 @@ GROUP BY agent_id, logical_seq HAVING COUNT(*) > 1; ``` -Expected result: no rows. Any returned row means that one agent has multiple -events with the same `logical_seq` and the capture/storage path is not safe -for that case. +Expected result: no rows. + +Logical sequence numbers are ordering values, not causal edges. Use +`causal_parent_ids` and the ancestor query for dependency relationships. + +## 12. Confirm Phase 2 does not write snapshots yet + +Purpose: confirms the `snapshots` table exists for the schema contract while +snapshot creation remains intentionally deferred to Phase 3. + +```sql +SELECT COUNT(*) AS snapshots +FROM snapshots; +``` + +Expected result: usually `0` for this Phase 2 test database. Existing rows are +not an error; Phase 2 does not create or reconstruct snapshots. diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..f440409 --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,5 @@ +"""Phase 2 graph operations.""" + +from .graph import ancestors, assign_causal_parents, record_causal_event + +__all__ = ["ancestors", "assign_causal_parents", "record_causal_event"] diff --git a/core/graph.py b/core/graph.py new file mode 100644 index 0000000..e5a8da8 --- /dev/null +++ b/core/graph.py @@ -0,0 +1,66 @@ +"""Explicit causal-parent assignment and graph queries.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from sdk.events import AgentClock, Event, next_seq + + +def _get_event(log: Any, event_id: str) -> Event: + getter = getattr(log, "get", None) + event = getter(event_id) if getter is not None else None + if event is None: + raise ValueError(f"causal parent event {event_id} does not exist") + return event + + +def assign_causal_parents( + agent_id: str, + clock: AgentClock, + causal_parents: list[str], + log: Any, +) -> int: + """Validate explicit parents and allocate the dependent event sequence.""" + parent_events = [_get_event(log, event_id) for event_id in causal_parents] + parent_seqs = [event.logical_seq for event in parent_events] + allocator = getattr(log, "allocate_logical_seq", None) + if allocator is not None: + return allocator(agent_id, clock, parent_seqs) + return next_seq(clock, parent_seqs) + + +def record_causal_event( + *, + agent_id: str, + clock: AgentClock, + log: Any, + event_type: str, + payload: dict[str, Any], + causal_parents: Iterable[str], + idempotency_key: str | None = None, + run_id: str | None = None, +) -> Event: + """Append an event whose explicit parents were used by the caller.""" + parent_ids = list(causal_parents) + logical_seq = assign_causal_parents(agent_id, clock, parent_ids, log) + event = Event( + agent_id=agent_id, + logical_seq=logical_seq, + event_type=event_type, + payload=payload, + causal_parent_ids=parent_ids, + idempotency_key=idempotency_key, + run_id=run_id, + ) + result = log.append(event) + return result if isinstance(result, Event) else event + + +def ancestors(event_id: str, log: Any) -> set[str]: + """Return the starting event and all of its PostgreSQL-backed ancestors.""" + getter = getattr(log, "ancestors", None) + if getter is None: + raise TypeError("ancestors requires a storage adapter with an ancestors method") + return set(getter(event_id)) diff --git a/storage/postgres.py b/storage/postgres.py index 7057979..5fec574 100644 --- a/storage/postgres.py +++ b/storage/postgres.py @@ -1,7 +1,7 @@ -"""Minimal psycopg 3 storage adapter for Phase 1. +"""Minimal psycopg 3 storage adapter for Phase 1 and Phase 2. -The adapter intentionally owns only event and agent persistence. Graph -queries, reducers, snapshots, and provenance are later phases. +The adapter owns the Phase 1 event/agent writes and the Phase 2 graph query. +State reconstruction, slicing, and provenance remain later phases. """ from __future__ import annotations @@ -55,12 +55,24 @@ created_at timestamptz NOT NULL DEFAULT now(), UNIQUE (agent_id, logical_seq) ); +CREATE TABLE IF NOT EXISTS snapshots ( + id uuid PRIMARY KEY DEFAULT gen_random_uuid(), + run_id uuid NOT NULL REFERENCES runs(id), + agent_id uuid NOT NULL REFERENCES agents(id), + logical_seq bigint NOT NULL, + state jsonb NOT NULL, + state_hash text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now() +); ALTER TABLE agents DROP CONSTRAINT IF EXISTS fk_spawned_at_event; ALTER TABLE agents ADD CONSTRAINT fk_spawned_at_event FOREIGN KEY (spawned_at_event_id) REFERENCES events(id); CREATE UNIQUE INDEX IF NOT EXISTS idx_events_idempotency ON events (agent_id, idempotency_key) WHERE idempotency_key IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_events_agent_seq ON events (agent_id, logical_seq); +CREATE INDEX IF NOT EXISTS idx_events_run_seq ON events (run_id, logical_seq); +CREATE INDEX IF NOT EXISTS idx_snapshots_agent_seq + ON snapshots (agent_id, logical_seq DESC); """ @@ -190,6 +202,34 @@ def get_by_idempotency_key(self, agent_id: str, key: str) -> Event | None: row = cursor.fetchone() return None if row is None else self._row_to_event(row) + def ancestors(self, event_id: str) -> list[str]: + """Return an event and all events reachable through causal parents.""" + event_uuid = self._uuid(event_id, "event_id") + with self.connection.cursor() as cursor: + cursor.execute( + """ + WITH RECURSIVE ancestors AS ( + SELECT id, causal_parent_ids + FROM events + WHERE id = %s + + UNION + + SELECT e.id, e.causal_parent_ids + FROM events e + JOIN ancestors a ON e.id = ANY(a.causal_parent_ids) + ) + SELECT id::text + FROM ancestors + ORDER BY id::text + """, + (event_uuid,), + ) + rows = cursor.fetchall() + if not rows: + raise ValueError(f"event {event_id} does not exist") + return [str(row[0]) for row in rows] + @contextmanager def tool_invocation_lock(self, agent_id: str, invocation_id: str) -> Iterator[None]: """Serialize one invocation across processes and database connections.""" diff --git a/tests/test_phase2.py b/tests/test_phase2.py new file mode 100644 index 0000000..dc6ebfb --- /dev/null +++ b/tests/test_phase2.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +import json +import os +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any +from uuid import UUID, uuid4 + +import pytest + +from core.graph import ancestors, assign_causal_parents, record_causal_event +from sdk.client import CapturedClient +from sdk.events import AgentClock, Event, InMemoryEventLog, record_event +from sdk.lifecycle import spawn_agent +from sdk.tools import capture_tool +from storage.postgres import PostgresEventStore + + +def test_assign_causal_parents_uses_parent_sequences_not_wall_time() -> None: + log = InMemoryEventLog() + log.append( + Event( + id="research-result", + agent_id="researcher", + logical_seq=4, + event_type="tool_result", + payload={}, + wall_time=datetime.now(UTC) + timedelta(days=1), + ) + ) + log.append( + Event( + id="coder-result", + agent_id="coder", + logical_seq=9, + event_type="tool_result", + payload={}, + wall_time=datetime.now(UTC) - timedelta(days=1), + ) + ) + + clock = AgentClock(counter=2) + sequence = assign_causal_parents( + "planner", + clock, + ["research-result", "coder-result"], + log, + ) + + assert sequence == 10 + assert clock.counter == 10 + + +def test_record_causal_event_preserves_multiple_parent_ids() -> None: + log = InMemoryEventLog() + parent_ids = ["parent-one", "parent-two"] + for parent_id in parent_ids: + log.append( + Event( + id=parent_id, + agent_id="worker", + logical_seq=1, + event_type="tool_result", + payload={}, + ) + ) + + event = record_causal_event( + agent_id="planner", + clock=AgentClock(), + log=log, + event_type="model_call", + payload={"action": "merge"}, + causal_parents=parent_ids, + ) + + assert event.causal_parent_ids == parent_ids + assert event.logical_seq == 2 + + +def test_assign_causal_parents_rejects_missing_parent() -> None: + with pytest.raises(ValueError, match="does not exist"): + assign_causal_parents("planner", AgentClock(), ["missing"], InMemoryEventLog()) + + +class _Response: + def model_dump(self) -> dict[str, str]: + return {"content": "ready"} + + +class _Messages: + def create(self, **_: object) -> _Response: + return _Response() + + +class _Anthropic: + def __init__(self) -> None: + self.messages = _Messages() + + +@pytest.mark.integration +def test_phase2_schema_and_real_three_agent_graph() -> None: + database_url = os.environ.get("DATABASE_URL") + if not database_url: + pytest.skip("set DATABASE_URL to run the real PostgreSQL integration test") + psycopg = pytest.importorskip("psycopg") + + run_id = str(uuid4()) + planner_id = str(uuid4()) + with psycopg.connect(database_url) as connection: + store = PostgresEventStore(connection, lock_dsn=database_url) + store.create_schema() + _assert_phase2_schema(connection) + + with connection.cursor() as cursor: + cursor.execute("INSERT INTO runs (id, name) VALUES (%s, %s)", (run_id, "phase-2")) + cursor.execute( + "INSERT INTO agents (id, run_id, role) VALUES (%s, %s, %s)", + (planner_id, run_id, "planner"), + ) + connection.commit() + + planner_clock = AgentClock() + researcher, researcher_spawn, researcher_clock = spawn_agent( + parent_agent_id=planner_id, + parent_clock=planner_clock, + run_id=run_id, + role="researcher", + log=store, + agent_store=store, + child_agent_id=str(uuid4()), + ) + coder, coder_spawn, coder_clock = spawn_agent( + parent_agent_id=planner_id, + parent_clock=planner_clock, + run_id=run_id, + role="coder", + log=store, + agent_store=store, + child_agent_id=str(uuid4()), + ) + + researcher_client = CapturedClient( + None, + researcher.id, + researcher_clock, + store, + client=_Anthropic(), + run_id=run_id, + ) + coder_client = CapturedClient( + None, + coder.id, + coder_clock, + store, + client=_Anthropic(), + run_id=run_id, + ) + researcher_model = researcher_client.messages.create( + model="test-model", + max_tokens=10, + messages=[{"role": "user", "content": "research"}], + causal_parent_ids=[researcher_spawn.id], + ) + coder_model = coder_client.messages.create( + model="test-model", + max_tokens=10, + messages=[{"role": "user", "content": "code"}], + causal_parent_ids=[coder_spawn.id], + ) + + @capture_tool + def inspect(label: str) -> dict[str, str]: + return {label: "done"} + + inspect( + "research", + agent_id=researcher.id, + clock=researcher_clock, + log=store, + run_id=run_id, + causal_parent_ids=[_latest_event_id(store, researcher.id)], + invocation_id="phase2-research-tool", + ) + inspect( + "code", + agent_id=coder.id, + clock=coder_clock, + log=store, + run_id=run_id, + causal_parent_ids=[_latest_event_id(store, coder.id)], + invocation_id="phase2-code-tool", + ) + researcher_result = _latest_event(store, researcher.id, "tool_result") + coder_result = _latest_event(store, coder.id, "tool_result") + + merge = record_causal_event( + agent_id=planner_id, + clock=planner_clock, + log=store, + event_type="model_call", + payload={ + "model": "test-model", + "input": [researcher_model.model_dump(), coder_model.model_dump()], + "output": {"merged": True}, + }, + causal_parents=[researcher_result.id, coder_result.id], + run_id=run_id, + ) + + researcher_record = store.get_agent(researcher.id) + coder_record = store.get_agent(coder.id) + assert researcher_record is not None + assert coder_record is not None + assert researcher_record.parent_agent_id == planner_id + assert coder_record.parent_agent_id == planner_id + assert researcher_record.spawned_at_event_id == researcher_spawn.id + assert coder_record.spawned_at_event_id == coder_spawn.id + assert merge.causal_parent_ids == [researcher_result.id, coder_result.id] + + fixture = json.loads((Path(__file__).parents[1] / "fixture" / "fixture.json").read_text()) + fixture_agents = {agent["id"]: agent for agent in fixture["agents"]} + fixture_roles = {agent["role"] for agent in fixture["agents"]} + assert {"planner", "researcher", "coder"} <= fixture_roles + planner_record = store.get_agent(planner_id) + assert planner_record is not None + assert { + planner_record.role, + researcher_record.role, + coder_record.role, + } == {"planner", "researcher", "coder"} + for role, agent in (("researcher", researcher), ("coder", coder)): + fixture_agent_id = next( + agent_id for agent_id, value in fixture_agents.items() if value["role"] == role + ) + fixture_branch = [ + event["event_type"] + for event in fixture["events"] + if event["agent_id"] == fixture_agent_id + ] + actual_branch = [ + event.event_type + for event_id in _event_ids(store, agent.id) + if (event := store.get(event_id)) is not None + ] + assert actual_branch == fixture_branch + fixture_merge = next( + event + for event in fixture["events"] + if event["agent_id"] == "A" and len(event["causal_parent_ids"]) == 3 + ) + fixture_events = {event["id"]: event for event in fixture["events"]} + fixture_cross_parent_count = sum( + fixture_events[parent_id]["agent_id"] != "A" + for parent_id in fixture_merge["causal_parent_ids"] + ) + assert len(merge.causal_parent_ids) == fixture_cross_parent_count + researcher_branch = ancestors(researcher_result.id, store) + coder_branch = ancestors(coder_result.id, store) + assert ancestors(merge.id, store) == { + merge.id, + *researcher_branch, + *coder_branch, + } + assert researcher_branch.isdisjoint(coder_branch) + + root = record_event( + agent_id=planner_id, + clock=planner_clock, + log=store, + event_type="context_update", + payload={"shared": True}, + run_id=run_id, + ) + left = record_causal_event( + agent_id=researcher.id, + clock=researcher_clock, + log=store, + event_type="context_update", + payload={"branch": "left"}, + causal_parents=[root.id], + run_id=run_id, + ) + right = record_causal_event( + agent_id=coder.id, + clock=coder_clock, + log=store, + event_type="context_update", + payload={"branch": "right"}, + causal_parents=[root.id], + run_id=run_id, + ) + shared_merge = record_causal_event( + agent_id=planner_id, + clock=planner_clock, + log=store, + event_type="context_update", + payload={"shared_merge": True}, + causal_parents=[left.id, right.id], + run_id=run_id, + ) + assert ancestors(shared_merge.id, store) == { + shared_merge.id, + left.id, + right.id, + root.id, + } + assert ancestors(root.id, store) == {root.id} + with pytest.raises(ValueError, match="does not exist"): + store.ancestors(str(uuid4())) + + +def _latest_event(store: PostgresEventStore, agent_id: str, event_type: str): + event_ids = _event_ids(store, agent_id) + events = [store.get(event_id) for event_id in event_ids] + return next(event for event in reversed(events) if event and event.event_type == event_type) + + +def _latest_event_id(store: PostgresEventStore, agent_id: str) -> str: + return _event_ids(store, agent_id)[-1] + + +def _event_ids(store: PostgresEventStore, agent_id: str) -> list[str]: + with store.connection.cursor() as cursor: + cursor.execute( + "SELECT id FROM events WHERE agent_id = %s ORDER BY logical_seq", + (UUID(agent_id),), + ) + return [str(row[0]) for row in cursor.fetchall()] + + +def _assert_phase2_schema(connection: Any) -> None: + with connection.cursor() as cursor: + cursor.execute( + "SELECT table_name FROM information_schema.tables " + "WHERE table_schema = 'public' AND table_name IN " + "('runs', 'agents', 'events', 'snapshots')" + ) + assert {row[0] for row in cursor.fetchall()} == { + "runs", + "agents", + "events", + "snapshots", + } + + cursor.execute( + "SELECT table_name, column_name, udt_name " + "FROM information_schema.columns " + "WHERE table_schema = 'public' AND table_name IN " + "('runs', 'agents', 'events', 'snapshots')" + ) + columns = {(row[0], row[1]): row[2] for row in cursor.fetchall()} + required_columns = { + ("runs", "id"), + ("agents", "run_id"), + ("agents", "parent_agent_id"), + ("agents", "spawned_at_event_id"), + ("agents", "lamport_offset"), + ("events", "run_id"), + ("events", "agent_id"), + ("events", "logical_seq"), + ("events", "wall_time"), + ("events", "event_type"), + ("events", "causal_parent_ids"), + ("events", "payload"), + ("events", "idempotency_key"), + ("snapshots", "run_id"), + ("snapshots", "agent_id"), + ("snapshots", "logical_seq"), + ("snapshots", "state"), + ("snapshots", "state_hash"), + } + assert required_columns <= columns.keys() + assert columns["events", "id"] == "uuid" + assert columns["events", "agent_id"] == "uuid" + assert columns["events", "causal_parent_ids"] == "_uuid" + assert columns["snapshots", "state"] == "jsonb" + assert columns["snapshots", "state_hash"] == "text" + + cursor.execute( + "SELECT indexname FROM pg_indexes WHERE schemaname = 'public' " + "AND indexname IN ('idx_events_agent_seq', 'idx_events_run_seq', " + "'idx_snapshots_agent_seq', 'idx_events_idempotency')" + ) + assert {row[0] for row in cursor.fetchall()} == { + "idx_events_agent_seq", + "idx_events_run_seq", + "idx_snapshots_agent_seq", + "idx_events_idempotency", + } + + cursor.execute( + "SELECT tc.constraint_name, kcu.column_name " + "FROM information_schema.table_constraints tc " + "JOIN information_schema.key_column_usage kcu " + "ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema " + "WHERE tc.table_schema = 'public' AND tc.table_name = 'events' " + "AND tc.constraint_type = 'UNIQUE'" + ) + unique_constraints: dict[str, set[str]] = {} + for constraint_name, column_name in cursor.fetchall(): + unique_constraints.setdefault(constraint_name, set()).add(column_name) + assert {"agent_id", "logical_seq"} in unique_constraints.values() + + cursor.execute( + "SELECT kcu.table_name, kcu.column_name, ccu.table_name, ccu.column_name " + "FROM information_schema.table_constraints tc " + "JOIN information_schema.key_column_usage kcu " + "ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema " + "JOIN information_schema.constraint_column_usage ccu " + "ON tc.constraint_name = ccu.constraint_name AND tc.table_schema = ccu.table_schema " + "WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_schema = 'public'" + ) + foreign_keys = {(row[0], row[1], row[2], row[3]) for row in cursor.fetchall()} + assert ("agents", "run_id", "runs", "id") in foreign_keys + assert ("agents", "parent_agent_id", "agents", "id") in foreign_keys + assert ("agents", "spawned_at_event_id", "events", "id") in foreign_keys + assert ("events", "run_id", "runs", "id") in foreign_keys + assert ("events", "agent_id", "agents", "id") in foreign_keys + assert ("snapshots", "run_id", "runs", "id") in foreign_keys + assert ("snapshots", "agent_id", "agents", "id") in foreign_keys From c39717e28d689100fa72cd0b35aa09282f263e93 Mon Sep 17 00:00:00 2001 From: DSCmatter Date: Tue, 18 Aug 2026 12:23:21 +0530 Subject: [PATCH 2/4] fix: various stuff --- core/graph.py | 13 ++- sdk/events.py | 6 ++ storage/postgres.py | 25 ++++- tests/test_phase1.py | 25 +++++ tests/test_phase2.py | 44 +++++++++ tests/test_postgres_integration.py | 145 ++++++++++++++++++++++++++++- 6 files changed, 251 insertions(+), 7 deletions(-) diff --git a/core/graph.py b/core/graph.py index e5a8da8..d07c02a 100644 --- a/core/graph.py +++ b/core/graph.py @@ -21,9 +21,14 @@ def assign_causal_parents( clock: AgentClock, causal_parents: list[str], log: Any, + run_id: str | None = None, ) -> int: """Validate explicit parents and allocate the dependent event sequence.""" parent_events = [_get_event(log, event_id) for event_id in causal_parents] + if run_id is not None: + for event in parent_events: + if event.run_id is not None and event.run_id != run_id: + raise ValueError(f"causal parent event {event.id} belongs to another run") parent_seqs = [event.logical_seq for event in parent_events] allocator = getattr(log, "allocate_logical_seq", None) if allocator is not None: @@ -44,7 +49,13 @@ def record_causal_event( ) -> Event: """Append an event whose explicit parents were used by the caller.""" parent_ids = list(causal_parents) - logical_seq = assign_causal_parents(agent_id, clock, parent_ids, log) + if idempotency_key is not None: + getter = getattr(log, "get_by_idempotency_key", None) + if getter is not None: + existing = getter(agent_id, idempotency_key) + if isinstance(existing, Event): + return existing + logical_seq = assign_causal_parents(agent_id, clock, parent_ids, log, run_id) event = Event( agent_id=agent_id, logical_seq=logical_seq, diff --git a/sdk/events.py b/sdk/events.py index 0689894..56096be 100644 --- a/sdk/events.py +++ b/sdk/events.py @@ -96,6 +96,12 @@ def record_event( ) -> Event: """Allocate and append an event while preserving its causal metadata.""" parent_ids = list(causal_parent_ids) + if idempotency_key is not None: + getter = getattr(log, "get_by_idempotency_key", None) + if getter is not None: + existing = getter(agent_id, idempotency_key) + if isinstance(existing, Event): + return existing allocator = getattr(log, "allocate_logical_seq", None) if allocator is None: sequence = next_seq(clock, causal_parent_seqs) diff --git a/storage/postgres.py b/storage/postgres.py index 5fec574..9759298 100644 --- a/storage/postgres.py +++ b/storage/postgres.py @@ -112,6 +112,15 @@ def append(self, event: Event) -> Event: """ try: with self.connection.cursor() as cursor: + if parent_ids: + cursor.execute( + "SELECT id FROM events " + "WHERE id = ANY(%s) AND run_id = %s", + (parent_ids, run_id), + ) + owned_parent_ids = {row[0] for row in cursor.fetchall()} + if owned_parent_ids != set(parent_ids): + raise ValueError("causal parent events must belong to the event run") cursor.execute( sql, ( @@ -135,7 +144,12 @@ def append(self, event: Event) -> Event: (agent_id, event.idempotency_key), ) row = cursor.fetchone() - self.connection.commit() + if row is not None: + # Roll back a sequence allocated in this transaction + # when another writer already stored this idempotent event. + self.connection.rollback() + if row is not None: + self.connection.commit() except Exception: self.connection.rollback() raise @@ -172,7 +186,6 @@ def allocate_logical_seq( "UPDATE agents SET lamport_offset = %s WHERE id = %s", (sequence, agent_uuid), ) - self.connection.commit() except Exception: self.connection.rollback() raise @@ -209,15 +222,17 @@ def ancestors(self, event_id: str) -> list[str]: cursor.execute( """ WITH RECURSIVE ancestors AS ( - SELECT id, causal_parent_ids + SELECT id, run_id, causal_parent_ids FROM events WHERE id = %s UNION - SELECT e.id, e.causal_parent_ids + SELECT e.id, e.run_id, e.causal_parent_ids FROM events e - JOIN ancestors a ON e.id = ANY(a.causal_parent_ids) + JOIN ancestors a + ON e.id = ANY(a.causal_parent_ids) + AND e.run_id = a.run_id ) SELECT id::text FROM ancestors diff --git a/tests/test_phase1.py b/tests/test_phase1.py index 774d64b..88ae212 100644 --- a/tests/test_phase1.py +++ b/tests/test_phase1.py @@ -45,6 +45,31 @@ def append(self, event: Event) -> Event: assert event.causal_parent_ids == ["parent-1", "parent-2"] +def test_record_event_idempotent_retry_does_not_allocate_again() -> None: + log = InMemoryEventLog() + clock = AgentClock() + first = record_event( + agent_id="a", + clock=clock, + log=log, + event_type="context_update", + payload={"attempt": 1}, + idempotency_key="same-event", + ) + + retry = record_event( + agent_id="a", + clock=clock, + log=log, + event_type="context_update", + payload={"attempt": 2}, + idempotency_key="same-event", + ) + + assert retry == first + assert clock.current() == first.logical_seq + + @dataclass class FakeResponse: answer: str diff --git a/tests/test_phase2.py b/tests/test_phase2.py index dc6ebfb..9f4b056 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -79,11 +79,55 @@ def test_record_causal_event_preserves_multiple_parent_ids() -> None: assert event.logical_seq == 2 +def test_record_causal_event_idempotent_retry_does_not_allocate_again() -> None: + log = InMemoryEventLog() + clock = AgentClock() + first = record_causal_event( + agent_id="planner", + clock=clock, + log=log, + event_type="context_update", + payload={"attempt": 1}, + causal_parents=[], + idempotency_key="same-merge", + ) + + retry = record_causal_event( + agent_id="planner", + clock=clock, + log=log, + event_type="context_update", + payload={"attempt": 2}, + causal_parents=[], + idempotency_key="same-merge", + ) + + assert retry == first + assert clock.current() == first.logical_seq + + def test_assign_causal_parents_rejects_missing_parent() -> None: with pytest.raises(ValueError, match="does not exist"): assign_causal_parents("planner", AgentClock(), ["missing"], InMemoryEventLog()) +def test_assign_causal_parents_rejects_cross_run_parent() -> None: + log = InMemoryEventLog() + log.append( + Event( + id="parent", + run_id="run-a", + agent_id="worker", + logical_seq=1, + event_type="tool_result", + payload={}, + ) + ) + + with pytest.raises(ValueError, match="another run"): + assign_causal_parents("planner", AgentClock(), ["parent"], log, "run-b") + + class _Response: def model_dump(self) -> dict[str, str]: return {"content": "ready"} diff --git a/tests/test_postgres_integration.py b/tests/test_postgres_integration.py index 011a17c..007c69f 100644 --- a/tests/test_postgres_integration.py +++ b/tests/test_postgres_integration.py @@ -8,8 +8,9 @@ import pytest +from core.graph import record_causal_event from sdk.client import CapturedClient -from sdk.events import AgentClock, record_event +from sdk.events import AgentClock, Event, record_event from sdk.lifecycle import spawn_agent from sdk.tools import capture_tool from storage.postgres import PostgresEventStore @@ -180,6 +181,148 @@ def write_event(store: PostgresEventStore, clock: AgentClock) -> int: assert sorted(sequences) == [1, 2] +@pytest.mark.integration +def test_postgres_sequence_allocation_rolls_back_with_failed_append() -> None: + database_url = os.environ.get("DATABASE_URL") + if not database_url: + pytest.skip("set DATABASE_URL to run the real Postgres integration test") + psycopg = pytest.importorskip("psycopg") + run_id = str(uuid4()) + agent_id = str(uuid4()) + with psycopg.connect(database_url) as connection: + store = PostgresEventStore(connection, lock_dsn=database_url) + store.create_schema() + with connection.cursor() as cursor: + cursor.execute("INSERT INTO runs (id, name) VALUES (%s, %s)", (run_id, "sequence")) + cursor.execute( + "INSERT INTO agents (id, run_id, role) VALUES (%s, %s, %s)", + (agent_id, run_id, "worker"), + ) + connection.commit() + + with pytest.raises(psycopg.Error): + record_event( + agent_id=agent_id, + clock=AgentClock(), + log=store, + event_type="not_an_event_type", + payload={}, + run_id=run_id, + ) + + event = record_event( + agent_id=agent_id, + clock=AgentClock(), + log=store, + event_type="context_update", + payload={}, + idempotency_key="sequence-retry", + run_id=run_id, + ) + retry = record_event( + agent_id=agent_id, + clock=AgentClock(), + log=store, + event_type="context_update", + payload={"changed": True}, + idempotency_key="sequence-retry", + run_id=run_id, + ) + + assert event.logical_seq == 1 + assert retry.id == event.id + assert retry.logical_seq == 1 + with connection.cursor() as cursor: + cursor.execute("SELECT lamport_offset FROM agents WHERE id = %s", (agent_id,)) + assert cursor.fetchone()[0] == 1 + + +@pytest.mark.integration +def test_postgres_graph_rejects_and_isolates_cross_run_parents() -> None: + database_url = os.environ.get("DATABASE_URL") + if not database_url: + pytest.skip("set DATABASE_URL to run the real Postgres integration test") + psycopg = pytest.importorskip("psycopg") + from psycopg.types.json import Jsonb + + run_a, run_b = str(uuid4()), str(uuid4()) + agent_a, agent_b = str(uuid4()), str(uuid4()) + with psycopg.connect(database_url) as connection: + store = PostgresEventStore(connection, lock_dsn=database_url) + store.create_schema() + with connection.cursor() as cursor: + cursor.execute("INSERT INTO runs (id, name) VALUES (%s, %s)", (run_a, "graph-a")) + cursor.execute("INSERT INTO runs (id, name) VALUES (%s, %s)", (run_b, "graph-b")) + cursor.execute( + "INSERT INTO agents (id, run_id, role) VALUES (%s, %s, %s)", + (agent_a, run_a, "worker-a"), + ) + cursor.execute( + "INSERT INTO agents (id, run_id, role) VALUES (%s, %s, %s)", + (agent_b, run_b, "worker-b"), + ) + connection.commit() + + parent = record_event( + agent_id=agent_a, + clock=AgentClock(), + log=store, + event_type="context_update", + payload={}, + run_id=run_a, + ) + with pytest.raises(ValueError, match="another run"): + record_causal_event( + agent_id=agent_b, + clock=AgentClock(), + log=store, + event_type="context_update", + payload={}, + causal_parents=[parent.id], + run_id=run_b, + ) + + with pytest.raises(ValueError, match="must belong to the event run"): + store.append( + Event( + run_id=run_b, + agent_id=agent_b, + logical_seq=1, + event_type="context_update", + payload={}, + causal_parent_ids=[parent.id], + ) + ) + + legacy_event = Event( + run_id=run_b, + agent_id=agent_b, + logical_seq=1, + event_type="context_update", + payload={}, + causal_parent_ids=[parent.id], + ) + record = legacy_event.to_record() + with connection.cursor() as cursor: + cursor.execute( + "INSERT INTO events " + "(id, run_id, agent_id, logical_seq, event_type, " + "causal_parent_ids, payload) VALUES (%s, %s, %s, %s, %s, %s, %s)", + ( + record["id"], + record["run_id"], + record["agent_id"], + record["logical_seq"], + record["event_type"], + record["causal_parent_ids"], + Jsonb(record["payload"]), + ), + ) + connection.commit() + + assert store.ancestors(legacy_event.id) == [legacy_event.id] + + def _event_ids(connection: Any, agent_id: str) -> list[str]: with connection.cursor() as cursor: cursor.execute( From e8efdeb32b296bbb4d0be86d76a96b7017a67b70 Mon Sep 17 00:00:00 2001 From: DSCmatter Date: Wed, 19 Aug 2026 19:27:48 +0530 Subject: [PATCH 3/4] fix --- core/graph.py | 2 ++ sdk/events.py | 2 ++ storage/postgres.py | 6 ++++-- tests/test_phase2.py | 23 +++++++++++++++++++++++ tests/test_postgres_integration.py | 14 ++++++++++++++ 5 files changed, 45 insertions(+), 2 deletions(-) diff --git a/core/graph.py b/core/graph.py index d07c02a..f083d53 100644 --- a/core/graph.py +++ b/core/graph.py @@ -55,6 +55,8 @@ def record_causal_event( existing = getter(agent_id, idempotency_key) if isinstance(existing, Event): return existing + if len(parent_ids) != len(set(parent_ids)): + raise ValueError("causal parent IDs must be unique") logical_seq = assign_causal_parents(agent_id, clock, parent_ids, log, run_id) event = Event( agent_id=agent_id, diff --git a/sdk/events.py b/sdk/events.py index 56096be..2b26473 100644 --- a/sdk/events.py +++ b/sdk/events.py @@ -102,6 +102,8 @@ def record_event( existing = getter(agent_id, idempotency_key) if isinstance(existing, Event): return existing + if len(parent_ids) != len(set(parent_ids)): + raise ValueError("causal parent IDs must be unique") allocator = getattr(log, "allocate_logical_seq", None) if allocator is None: sequence = next_seq(clock, causal_parent_seqs) diff --git a/storage/postgres.py b/storage/postgres.py index 9759298..1a74dd4 100644 --- a/storage/postgres.py +++ b/storage/postgres.py @@ -98,6 +98,8 @@ def append(self, event: Event) -> Event: parent_ids = [ self._uuid(value, "Event.causal_parent_ids") for value in record["causal_parent_ids"] ] + if len(parent_ids) != len(set(parent_ids)): + raise ValueError("causal parent IDs must be unique") columns = ( "id, run_id, agent_id, logical_seq, wall_time, event_type, " "causal_parent_ids, payload, idempotency_key" @@ -148,13 +150,13 @@ def append(self, event: Event) -> Event: # Roll back a sequence allocated in this transaction # when another writer already stored this idempotent event. self.connection.rollback() - if row is not None: - self.connection.commit() + return self._row_to_event(row) except Exception: self.connection.rollback() raise if row is None: raise RuntimeError("event insert did not return an event") + self.connection.commit() return self._row_to_event(row) def allocate_logical_seq( diff --git a/tests/test_phase2.py b/tests/test_phase2.py index 9f4b056..cc13f73 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -106,6 +106,29 @@ def test_record_causal_event_idempotent_retry_does_not_allocate_again() -> None: assert clock.current() == first.logical_seq +def test_record_causal_event_rejects_duplicate_parent_ids() -> None: + log = InMemoryEventLog() + log.append( + Event( + id="parent", + agent_id="worker", + logical_seq=1, + event_type="tool_result", + payload={}, + ) + ) + + with pytest.raises(ValueError, match="must be unique"): + record_causal_event( + agent_id="planner", + clock=AgentClock(), + log=log, + event_type="context_update", + payload={}, + causal_parents=["parent", "parent"], + ) + + def test_assign_causal_parents_rejects_missing_parent() -> None: with pytest.raises(ValueError, match="does not exist"): assign_causal_parents("planner", AgentClock(), ["missing"], InMemoryEventLog()) diff --git a/tests/test_postgres_integration.py b/tests/test_postgres_integration.py index 007c69f..30153e8 100644 --- a/tests/test_postgres_integration.py +++ b/tests/test_postgres_integration.py @@ -219,6 +219,17 @@ def test_postgres_sequence_allocation_rolls_back_with_failed_append() -> None: idempotency_key="sequence-retry", run_id=run_id, ) + allocated = store.allocate_logical_seq(agent_id, AgentClock()) + conflicting_retry = store.append( + Event( + run_id=run_id, + agent_id=agent_id, + logical_seq=allocated, + event_type="context_update", + payload={"conflict": True}, + idempotency_key="sequence-retry", + ) + ) retry = record_event( agent_id=agent_id, clock=AgentClock(), @@ -230,6 +241,9 @@ def test_postgres_sequence_allocation_rolls_back_with_failed_append() -> None: ) assert event.logical_seq == 1 + assert allocated == 2 + assert conflicting_retry.id == event.id + assert conflicting_retry.logical_seq == 1 assert retry.id == event.id assert retry.logical_seq == 1 with connection.cursor() as cursor: From b00a7194cda8250f4ef21e160da020896145b584 Mon Sep 17 00:00:00 2001 From: DSCmatter Date: Wed, 19 Aug 2026 19:57:04 +0530 Subject: [PATCH 4/4] fix1 --- core/graph.py | 2 ++ storage/postgres.py | 23 ++++++++++++++++++----- tests/test_phase2.py | 5 +++++ tests/test_postgres_integration.py | 1 + 4 files changed, 26 insertions(+), 5 deletions(-) diff --git a/core/graph.py b/core/graph.py index f083d53..7f8b4c9 100644 --- a/core/graph.py +++ b/core/graph.py @@ -24,6 +24,8 @@ def assign_causal_parents( run_id: str | None = None, ) -> int: """Validate explicit parents and allocate the dependent event sequence.""" + if len(causal_parents) != len(set(causal_parents)): + raise ValueError("causal parent IDs must be unique") parent_events = [_get_event(log, event_id) for event_id in causal_parents] if run_id is not None: for event in parent_events: diff --git a/storage/postgres.py b/storage/postgres.py index 1a74dd4..3aa3c1b 100644 --- a/storage/postgres.py +++ b/storage/postgres.py @@ -95,11 +95,6 @@ def append(self, event: Event) -> Event: event_id = self._uuid(record["id"], "Event.id") run_id = self._uuid(record["run_id"], "Event.run_id") agent_id = self._uuid(record["agent_id"], "Event.agent_id") - parent_ids = [ - self._uuid(value, "Event.causal_parent_ids") for value in record["causal_parent_ids"] - ] - if len(parent_ids) != len(set(parent_ids)): - raise ValueError("causal parent IDs must be unique") columns = ( "id, run_id, agent_id, logical_seq, wall_time, event_type, " "causal_parent_ids, payload, idempotency_key" @@ -114,6 +109,24 @@ def append(self, event: Event) -> Event: """ try: with self.connection.cursor() as cursor: + if event.idempotency_key is not None: + cursor.execute( + "SELECT id, run_id, agent_id, logical_seq, wall_time, event_type, " + "causal_parent_ids, payload, idempotency_key FROM events " + "WHERE agent_id = %s AND idempotency_key = %s", + (agent_id, event.idempotency_key), + ) + existing_row = cursor.fetchone() + if existing_row is not None: + self.connection.rollback() + return self._row_to_event(existing_row) + + parent_ids = [ + self._uuid(value, "Event.causal_parent_ids") + for value in record["causal_parent_ids"] + ] + if len(parent_ids) != len(set(parent_ids)): + raise ValueError("causal parent IDs must be unique") if parent_ids: cursor.execute( "SELECT id FROM events " diff --git a/tests/test_phase2.py b/tests/test_phase2.py index cc13f73..5b34faf 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -129,6 +129,11 @@ def test_record_causal_event_rejects_duplicate_parent_ids() -> None: ) +def test_assign_causal_parents_rejects_duplicate_parent_ids() -> None: + with pytest.raises(ValueError, match="must be unique"): + assign_causal_parents("planner", AgentClock(), ["parent", "parent"], InMemoryEventLog()) + + def test_assign_causal_parents_rejects_missing_parent() -> None: with pytest.raises(ValueError, match="does not exist"): assign_causal_parents("planner", AgentClock(), ["missing"], InMemoryEventLog()) diff --git a/tests/test_postgres_integration.py b/tests/test_postgres_integration.py index 30153e8..2105d63 100644 --- a/tests/test_postgres_integration.py +++ b/tests/test_postgres_integration.py @@ -227,6 +227,7 @@ def test_postgres_sequence_allocation_rolls_back_with_failed_append() -> None: logical_seq=allocated, event_type="context_update", payload={"conflict": True}, + causal_parent_ids=["not-a-uuid"], idempotency_key="sequence-retry", ) )