perf: two durable transactions per agent step by ksarangmath · Pull Request #700 · omnara-ai/omnara · GitHub
Skip to content

perf: two durable transactions per agent step - #700

Open
ksarangmath wants to merge 17 commits into
fix/worker-graceful-drainfrom
perf/two-transaction-steps
Open

ksarangmath wants to merge 17 commits into
fix/worker-graceful-drainfrom
perf/two-transaction-steps

Conversation

@ksarangmath

@ksarangmath ksarangmath commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #680. It builds on #680's worker drain and shutdown lifecycle, so review and merge #680 first; GitHub will then retarget this PR to main.

The problem

Every agent step was a series of small transactions. Claim, model preparation, output acceptance, tool authorization, tool completion and lease release each committed separately. Each one re-locked the agent, re-checked the lease and reloaded execution state, and each write commit needs its own WAL flush. On hardware where the flush is slow, the flush rate was the throughput ceiling.

The change

An ordinary model step is now two durable transactions:

  1. Claim + intent. The claim admits work and persists the started model attempt in the same transaction. It returns a prepared-work descriptor that the executor consumes directly: configuration, revision, opening and watermark, all captured under the claim's locks. Provider setup and request encoding stay outside the transaction.
  2. Result + release. Accepted output, events, execution state, the webhook outbox and lease release commit together. Release happens before readiness publication, so the next step becomes claimable atomically with the result. Unit.Commit keeps its ordering: publication, outbox, resource checks, commit, notifications.

Tools that complete entirely in the database now declare that as a capability at registration, instead of being special-cased by name. Authorization, completion and batch release share one transaction for these tools: list_agents, read_agent, send_agent_message, list_processes, list_machines, inspect_machine, tool_search, list_interaction_handlers and set_interaction_handler. Tools with background effects (create_machine, delete_machine, spawn_agent, stop_agent) use a separate capability and keep their existing authorization path and resource-lock order.

What stays the same:

  • Fencing. Agents are locked in canonical order, then wakeup and runtime rows, then authority is validated against statement_timestamp(). Renewal still locks only the runtime row. No stale worker can commit after losing its lease.
  • Lost acknowledgements. A retry resolves the already-committed result by its immutable operation identity before it requires a live lease, for model output and for command-backed tools, including committed ownership transfers and replays that overlap resource revocation or deletion. A missing lease never means success, and resolving an old result never releases a newer claim or re-enqueues webhooks.
  • Unknown commit outcomes stay distinct from aborts. An uncertain claim never starts provider work, and outcome resolution is bounded by fix(worker): drain in-flight agent work on SIGTERM with a bounded shutdown #680's cleanup lifetime.
  • Worker cleanup. The worker gets an explicit step disposition (finished, retained or transferred) instead of always calling release. Renewal failures are coordinated with that disposition, so a step that committed is never cancelled by its own lease release. Genuine lease loss still cancels runtime-owned async work promptly.
  • Batches. Release is a batch decision. Deferred, permission-waiting, runtime-retained and transferred calls keep their semantics.
  • Recovery. Generic releaseRuntime recovery, retries, compaction handoffs and resume paths are unchanged.

Measurements

Isolated commit counts, asserted in a test:

Before After
Model step 4 2
Model step + list_agents 8 4

The load-test harness from #699 was run back to back on one laptop, against main and against this PR combined with #699 and the follow-up write-path PR. The laptop was thermally throttled for both runs, so the relative numbers matter here, not the absolute ones:

main combined
Explicit commits per model call 6.73 3.69
Realistic load (8,000 agents): model calls/s 74.9 109.2
Tool result → next model call, p50 85 s 43 s
Time to start a turn, p50 42 s 21 s
Connection-pool wait 38 ms 3 ms

Replay audit

Every path that can finish or transfer a claim resolves its persisted outcome before it requires the old lease or checks live eligibility:

Path How it completes Outcome it resolves by
Database-read tools (list_*, read_agent, inspect_machine, tool_search) Result, which can finish the batch Result ID and outcome
Runtime tools (files, web, skills, MCP and integration tools) Async result; the worker releases Result ID and outcome
run_command, upload_file, download_file Hands off to a process Process ID and command identity
read_process, write_process, stop_process Hands off to a process action, or returns an immediate result Action ID or result
create_machine, delete_machine, spawn_agent, send_agent_message, stop_agent Intent and result committed together Result plus binding or launch identity
set_interaction_handler Selection and result Result
ask_question, permission requests, answers, handler callbacks Hand off to an interaction, or resolve one Exact interaction ID; actor and callback authorization kept
Custom-tool callbacks, model output External result, or output plus finishing the claim Result, output and proposal identities

Each row has lost-acknowledgement and overlapping-replay tests, including preserving a newer lease.

Tests

New integration tests cover:

  • lost acknowledgements before and after commit, for claim+prepare, result+release and command-backed tools;
  • replay after a newer claim;
  • cancellation racing completion, in both lock orders;
  • renewal racing release, using the real renewal goroutine;
  • multi-tool and mixed async batches;
  • crashes at each boundary;
  • exactly-once outbox rows;
  • lifecycle lock order with production route fields, including the grant-revocation and machine-termination pairs;
  • shutdown during outcome resolution;
  • capability validation in the registry.

Mutation checks confirmed that removing automatic resolution, or restoring an early agent prelock, makes tests fail. Build, vet, lint and the affected packages' unit, integration and race tests pass locally. Tests that need blob storage run in CI.

🤖 Generated with Claude Code

ksarangmath and others added 14 commits October 19, 2026 13:15
…ng it

A worker deploy or scale-in canceled the root context on SIGTERM, so
in-flight model streams were dropped, their runtimes released as
runtime_released_before_model_result_acceptance, and another worker re-ran
those calls from scratch, paying for the same generation twice.

On the first SIGTERM or SIGINT the worker now stops claiming new work and
lets in-flight executions, their lease renewals, retained async tools,
worker control cancels and background tool tasks keep running until they
finish or OMNARA_WORKER_DRAIN_TIMEOUT passes. Then the remaining work is
canceled through the existing release path. A second signal cancels at
once, and 0s keeps the old immediate cancel.

The default is 25s so it fits inside the common 30s stop grace periods;
the orchestrator's stop timeout must stay above the drain timeout, or the
process is killed before it can release its leases. Compose now gives the
worker a 30s grace period instead of Docker's 10s default.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The drain deadline only canceled the worker's own executions. Background
tool tasks kept running past it, and a turn blocked in Submit on a full
background queue could hold shutdown until SIGKILL, before its runtime
lease was released. Cleanup after the deadline could also run for up to
a minute or more (lease/3 for release plus 30s retries for interrupted
async tool results), past any common stop timeout.

cmd/worker now owns a single deadline that cancels worker executions and
the background runner together, then gives cleanup a fixed 5s lifetime.
Release and async tool failure persistence derive from that lifetime, so
shutdown takes at most the drain timeout plus 5s plus process teardown.
The default drain drops to 20s so that fits the common 30s stop timeout,
and Compose's stop_grace_period becomes configurable.

The background runner now drains on Shutdown: it stops accepting, waits
for accepted and queued tasks until the shared deadline, and Submit
returns as soon as the runner stops accepting or is canceled. Claims
stop on the first signal, including a claim already waiting on the
database; a lease returned as that happens is still executed and
released.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…etime

Several cleanups on the worker's shutdown path ignored both cancellation
and the new cleanup lifetime. These were uploaded artifact blob deletion
(no timeout), the integration inbox failure finalizer (40s), inbox
artifact cleanup (30s), MCP refresh-lease ownership and release, model
failure notices, the limits catalog fetch and launch provisioning. Any
of them could keep the process alive past the documented stop timeout.

The lifetime helper moves to internal/cleanup so storage, integration,
mcp, kernel and machinepool can use it. cleanup.Detach keeps a context's
values and drops its cancellation, but ends at its own timeout or when
the shutdown cleanup lifetime ends. cmd/worker attaches that lifetime to
every component's context. Every detached cleanup reachable after the
first signal now ends within the drain timeout plus 5s.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Threading a cleanup lifetime through every component the worker waits on
did not converge. Database pool close and other teardown could still
outlive it, and the fixed timeouts it added shortened cleanup in normal
operation. This reverts those per-component changes, so artifact, inbox,
Discord, MCP, kernel, limits-catalog and provisioning cleanup behave as
on main.

Instead, cmd/worker arms a watchdog when shutdown starts, on the first
signal or a component failure. If the process is still running after
the drain timeout plus the 5s cleanup budget plus a 2s margin, it logs a
warning and exits non-zero. The documented stop-timeout requirement
(drain + 10s) now holds by construction; cleanup cut off by the watchdog
is best-effort, and unreleased leases expire normally.

The claim-race test keeps its exactly-once release assertion.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The watchdog logged its overrun warning synchronously before exiting, so a
blocked stdout could keep the process alive past the shutdown bound. Write
the warning from a goroutine, wait at most 100ms for it, and exit either way.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Return the durable attempt and captured configuration to the executor, retaining separate resume and compaction routes. Reuse acquisition, renewal, claim identity, and newly created attempt projections. The local monotonic lease budget still begins before acquisition; provider setup remains outside the transaction.
Release the exact fenced runtime before final readiness publication. Resolve committed results by attempt identity before requiring a live lease, and distinguish unknown outcomes from known aborts. Return explicit worker dispositions and defer renewal-loss cancellation while completion is being resolved. Preserve generic cleanup and separate retry and compaction transitions. Add deterministic lost-acknowledgement, crash, cancellation, renewal and outbox tests; update preparation and reaper expectations.
Merge permitted list_agents authorization with its transactional result and exact-lease finish. Retain deferred and runtime-owned work, transfer only batches with durable owners, and preserve every tool state notification. Resolve synchronous completed operations before fencing and verify exact result receipts after uncertain commits so replay never deletes a newer claim. Runtime-owned async completion keeps its original live-lease fence so fallback results cannot acknowledge a reaper result as their own.

Return tool projections from inserted rows, carry immutable root routing, and exercise cancellation, recovery, renewal, mixed ownership, and replay. A read-only pre-send status check prevents sequential descriptor replay and cancellation between claim and execution from issuing another provider request; the locked result fence remains authoritative.

Count the ordinary model plus list_agents protocol at eight explicit commits before and four after, with model-only work falling from four to two.
Reuse captured agent routes when commands reach their existing lock boundaries instead of prelocking agents before lifecycle, pool and machine resources. Recognize completed command dispatch before checking the deleted runtime row, and retain completion outcome and command identity checks on replay.

Share the caller cleanup lifetime with store resolution, and coordinate every exhausted renewal failure with pending completion. Genuine cancellation remains immediate. Add transaction and query fault seams at actual store boundaries, plus production-route lock races, lost acknowledgement and newer-claim tests, blocked resolution during worker shutdown, and a real renewal-goroutine timeout test.

The automatic-resolution tests fail if either store resolver is disabled. Restoring agent prelocking reproduces PostgreSQL 40P01 in the new grant-revocation barrier test.
…ecisions

Replace list_agents name checks with registered completion capabilities. Verify transactional-only phases and declared lock requirements for the nine bounded database tools; distinguish the four commands with background effects and keep their authorization and background dispatch paths intact. Require declared database completions to finish before the store commits.

Return an explicit authorization disposition to the kernel instead of inferring deferred authorization from an unchanged state. Cover registry rejection of incompatible phases or lock declarations and exercise a mixed batch through the executor and async scope, including a requeued call, runtime-owned async work, database completion and permission waiting.

The commit-count regression retains four to two commits for a model step and eight to four for a model plus list_agents cycle.
Resolve persisted waiting process ownership after lost acknowledgements and overlapping execution without releasing a newer claim. Defer integration and machine eligibility failures until committed operations are rechecked under canonical locks.

Add regression coverage for ownership receipts, unknown outcomes, overlapping waiting calls, integration revocation, and machine deletion.
@github-actions github-actions Bot added the perf PRs of type perf (auto-applied from title) label Oct 21, 2026

@cursor cursor Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

})
if err := unit.LockAgents(ctx, plan); err != nil {
_ = unit.Rollback(ctx)
return nil, nil, err

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Step lock misses inactive-lease mapping

Medium Severity

The new beginStepExecution fast path returns a raw ErrNotFound from LockAgents instead of mapping it to ErrRuntimeLockInactive the way lockExecution does when a runtimeID is present. Ordinary result+release now takes this path whenever RootAgentID is set and there is no parent effect, so a disappearing agent during accept or batch finish can look like a missing resource rather than a dead claim.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit b09b5e6. Configure here.

@ksarangmath
ksarangmath force-pushed the fix/worker-graceful-drain branch from d6381db to e3a439a Compare October 21, 2026 02:34

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 32e633b. Configure here.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale open interaction after unknown commit

Medium Severity

Unknown-commit resolution treats any persisted row with the same ID as success when the in-memory record is still open, even if the row is already resolved or canceled. CreatePermissionInteraction then returns that stale open record, so a permission prompt can be sent after the interaction has already been answered.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 32e633b. Configure here.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

perf PRs of type perf (auto-applied from title)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant