Repository navigation
feat: add phase 2 owner substrate - #79
Merged
Merged
Conversation
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Address review feedback around owner-store initialization, owner record typing, and test fixture protocol alignment. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds lease-based session ownership and fencing: a PostgreSQL-backed Changes
Sequence DiagramsequenceDiagram
participant Client
participant SessionManager
participant OwnerStore as SessionOwnerStore
participant Backend as PGSessionOwnerStore
participant DB as PostgreSQL<br/>session_owners
Client->>SessionManager: init(owner_id, fencing_token)
SessionManager->>SessionManager: validate owner metadata
rect rgba(100,150,200,0.5)
Note over Client,DB: Acquire ownership
Client->>SessionManager: run_agent(session_id)
SessionManager->>OwnerStore: acquire(session_id, owner_id, lease_seconds, fencing_token)
OwnerStore->>Backend: acquire(...)
Backend->>DB: UPSERT session_owners (if no unexpired lease)
DB-->>Backend: success / conflict
Backend-->>OwnerStore: bool
OwnerStore-->>SessionManager: bool
end
rect rgba(150,100,200,0.5)
Note over Client,DB: Ownership check
SessionManager->>OwnerStore: get_owner(session_id)
OwnerStore->>Backend: get_owner(session_id)
Backend->>DB: SELECT owner_id, lease_expires_at, fencing_token
DB-->>Backend: record (if lease unexpired)
Backend-->>OwnerStore: SessionOwnerRecord
OwnerStore-->>SessionManager: SessionOwnerRecord
SessionManager->>SessionManager: _assert_owner() validates match
end
rect rgba(200,150,100,0.5)
Note over Client,DB: Renew lease
SessionManager->>OwnerStore: renew(session_id, owner_id, lease_seconds, new_fencing_token, current_fencing_token)
OwnerStore->>Backend: renew(...)
Backend->>DB: UPDATE session_owners (if owner_id, token, lease match)
DB-->>Backend: success / conflict
Backend-->>OwnerStore: bool
OwnerStore-->>SessionManager: bool
end
rect rgba(200,100,150,0.5)
Note over Client,DB: Release ownership
Client->>SessionManager: close_session(session_id)
SessionManager->>SessionManager: _assert_owner() final check
SessionManager->>OwnerStore: release(session_id, owner_id, fencing_token)
OwnerStore->>Backend: release(...)
Backend->>DB: DELETE FROM session_owners (if owner_id & token match)
DB-->>Backend: success / conflict
Backend-->>OwnerStore: bool
OwnerStore-->>SessionManager: bool
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
coding-agent/tests/ui/test_session_manager_owner_checks.py (1)
19-35: Model live and expired leases separately in this fake.
FakeOwnerStore.acquire()writeslease_expires_at=datetime.now(UTC), so every acquired record is already expired at creation time. That makes the mismatch assertions pass, but it leaves this suite unable to distinguish a valid live lease from an expired one. Storenow + lease_secondshere and add one case whereowner_id/fencing_tokenstill match but the lease is in the past.Suggested adjustment
-from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta ... self._owners[session_id] = SessionOwnerRecord( owner_id=owner_id, - lease_expires_at=datetime.now(UTC), + lease_expires_at=datetime.now(UTC) + timedelta(seconds=lease_seconds), fencing_token=fencing_token, )Also applies to: 93-164
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@coding-agent/tests/ui/test_session_manager_owner_checks.py` around lines 19 - 35, The fake owner store's acquire method sets lease_expires_at to datetime.now(UTC) causing all leases to be immediately expired; change FakeOwnerStore.acquire to set lease_expires_at to datetime.now(UTC) + timedelta(seconds=lease_seconds) (use the SessionOwnerRecord fields owner_id, lease_expires_at, fencing_token) so live vs expired leases can be distinguished, and add one test case where owner_id and fencing_token match but lease_expires_at is in the past to verify expired-lease behavior; also apply the same lease_expires_at handling to the other fake owner-store methods in the file that create SessionOwnerRecord (lines referenced in the review) so they consistently use now + lease_seconds.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@coding-agent/src/coding_agent/ui/session_manager.py`:
- Around line 964-966: The ownership check in session_manager methods is
performed before awaiting the local lock, allowing a lease change to slip in
while blocked; move the call to _assert_owner(session_id) so it executes after
acquiring the lock (inside the async with self._lock or the corresponding
turn_lock scope) and immediately before any destructive operations (e.g., before
calling get_session_async, close_session logic, or run_agent's first mutating
step). Update the occurrences around the current calls at get_session_async
(lines ~964) and the similar block at ~1015-1020 to call await
self._assert_owner(session_id) inside the lock context to revalidate ownership
before proceeding.
- Around line 493-510: The _assert_owner method currently only checks owner_id
and fencing_token but must also reject expired leases: after fetching owner =
await self._owner_store.get_owner(session_id) inspect owner.lease_expires_at and
compare to the current time (e.g., datetime.utcnow() or timezone-aware
equivalent); if lease_expires_at is in the past (or missing/None if you treat
that as expired) raise a RuntimeError (e.g., "session lease expired") so stale
SessionOwnerRecord entries are not accepted—update _assert_owner to perform this
expiry check before validating owner_id/fencing_token (references:
_assert_owner, self._owner_store.get_owner, SessionOwnerRecord.lease_expires_at,
owner.owner_id, owner.fencing_token).
In `@coding-agent/tests/agentkit/storage/test_pg.py`:
- Around line 110-122: The FakePool test double is storing raw lease payloads
and returning expired rows unfiltered, diverging from the real
PGSessionOwnerStore behavior; update FakePool’s insert/update branches that
touch session_owners (the block handling "INSERT INTO session_owners" and the
similar blocks at 123-143 and 185-195) to compute and store a proper
lease_expires_at timestamp (e.g., current time + lease_seconds or a comparable
numeric expiry) instead of the raw third argument, and ensure the read path (the
logic used by get_owner in tests) filters out rows whose stored lease_expires_at
is expired (past timestamp) so FakePool enforces the same expiry suppression and
typing as PGSessionOwnerStore.
---
Nitpick comments:
In `@coding-agent/tests/ui/test_session_manager_owner_checks.py`:
- Around line 19-35: The fake owner store's acquire method sets lease_expires_at
to datetime.now(UTC) causing all leases to be immediately expired; change
FakeOwnerStore.acquire to set lease_expires_at to datetime.now(UTC) +
timedelta(seconds=lease_seconds) (use the SessionOwnerRecord fields owner_id,
lease_expires_at, fencing_token) so live vs expired leases can be distinguished,
and add one test case where owner_id and fencing_token match but
lease_expires_at is in the past to verify expired-lease behavior; also apply the
same lease_expires_at handling to the other fake owner-store methods in the file
that create SessionOwnerRecord (lines referenced in the review) so they
consistently use now + lease_seconds.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4a516de2-5a66-407e-a4d9-bc24c5db6148
📒 Files selected for processing (6)
coding-agent/src/agentkit/storage/pg.pycoding-agent/src/coding_agent/ui/session_manager.pycoding-agent/src/coding_agent/ui/session_owner_store.pycoding-agent/tests/agentkit/storage/test_pg.pycoding-agent/tests/ui/test_session_manager_owner_checks.pycoding-agent/tests/ui/test_session_owner_store.py
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@coding-agent/src/coding_agent/ui/session_manager.py`:
- Around line 385-391: The constructor currently only rejects when owner_store
is set but owner_id or fencing_token are missing; add the inverse validation to
also reject when owner_store is None but either owner_id or fencing_token is
provided. In the SessionManager constructor (the code handling owner_store,
owner_id, and fencing_token), raise a ValueError with a clear message like
"owner_store must be provided when owner_id or fencing_token is set" if
owner_store is None and (owner_id is not None or fencing_token is not None), so
ownership cannot be silently disabled.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 71057c1d-ed62-474f-9dbf-757b1782b21c
📒 Files selected for processing (3)
coding-agent/src/coding_agent/ui/session_manager.pycoding-agent/tests/agentkit/storage/test_pg.pycoding-agent/tests/ui/test_session_manager_owner_checks.py
✅ Files skipped from review due to trivial changes (1)
- coding-agent/tests/ui/test_session_manager_owner_checks.py
There was a problem hiding this comment.
Pull request overview
Adds the Phase 2 “owner substrate” for multi-instance safety by introducing a PostgreSQL-backed session ownership/lease store and applying owner/fencing validation to key session runtime operations.
Changes:
- Introduces
PGSessionOwnerStoreplus a typedSessionOwnerRecordwrapper (SessionOwnerStore) for validating/normalizing owner payloads. - Adds owner/fencing validation hooks to
SessionManagerruntime entry points (run/close/checkpoint/restore paths). - Adds targeted unit tests for the owner store and owner gating behavior.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| coding-agent/src/agentkit/storage/pg.py | Adds PGSessionOwnerStore with schema + acquire/renew/release/get operations backed by Postgres. |
| coding-agent/src/coding_agent/ui/session_owner_store.py | Adds typed owner record abstraction and a wrapper store that normalizes backend payloads to SessionOwnerRecord. |
| coding-agent/src/coding_agent/ui/session_manager.py | Adds owner/fencing parameters + _assert_owner() and integrates owner checks into several runtime entry points. |
| coding-agent/tests/agentkit/storage/test_pg.py | Extends PG fakes and adds coverage for PGSessionOwnerStore. |
| coding-agent/tests/ui/test_session_owner_store.py | Adds tests validating typed owner record behavior for the UI-layer owner store wrapper. |
| coding-agent/tests/ui/test_session_manager_owner_checks.py | Adds tests ensuring stale/expired ownership is rejected across key SessionManager operations. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+493
to
+513
| async def _assert_owner(self, session_id: str) -> None: | ||
| if self._owner_store is None: | ||
| return | ||
| if self._owner_id is None or self._fencing_token is None: | ||
| raise RuntimeError("stale owner or fencing token rejected") | ||
|
|
||
| owner = await self._owner_store.get_owner(session_id) | ||
| if owner is None: | ||
| raise RuntimeError("session has no owner") | ||
| if owner.lease_expires_at <= datetime.now(UTC): | ||
| raise RuntimeError("session owner lease expired") | ||
|
|
||
| current_owner_id = owner.owner_id | ||
| current_fencing_token = owner.fencing_token | ||
|
|
||
| if ( | ||
| current_owner_id != self._owner_id | ||
| or current_fencing_token != self._fencing_token | ||
| ): | ||
| raise RuntimeError("stale owner or fencing token rejected") | ||
|
|
There was a problem hiding this comment.
Owner/fencing enforcement is only applied on a subset of mutating entry points. shutdown_session_runtime() and ensure_session_runtime() (and also submit_approval() / wait_for_http_approval()) mutate/persist session state but currently do not call _assert_owner(), which allows a non-owner or stale instance to clear turn_in_progress, alter tape_id, or record approvals despite failing run_agent/checkpoint/close checks. Add await self._assert_owner(session_id) early in those methods (ideally under the same lock/turn_lock used for the mutation) to make fencing consistent.
| return _WireConsumer(session.wire, _request_approval) | ||
|
|
||
| async def _restore_checkpoint(self, session: Session, checkpoint_id: str) -> None: | ||
| await self._assert_owner(session.id) |
There was a problem hiding this comment.
restore_checkpoint() already calls _assert_owner(session_id) under the turn lock, and _restore_checkpoint() now calls _assert_owner(session.id) again. This duplicates an owner-store read and can also create a confusing failure mode where the first check passes but the second fails if the lease expires between them. Prefer enforcing ownership in exactly one place (either the public method before calling _restore_checkpoint, or inside _restore_checkpoint and remove the outer call).
Suggested change
| await self._assert_owner(session.id) |
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
coding-agent/src/coding_agent/ui/session_manager.py (2)
1166-1175:⚠️ Potential issue | 🟠 MajorPreserve the existing “Session not found” behavior here.
Calling
_assert_owner()before the existence check makes theApprovalResponse(..., feedback="Session not found")branch unreachable whenever fencing is enabled and the owner row is absent. Missing sessions now raiseRuntimeError("session has no owner")instead of returning the method’s current fallback result.Suggested fix
async def wait_for_http_approval( self, session_id: str, approval_req: ApprovalRequest, timeout_seconds: float, ) -> ApprovalResponse: - await self._assert_owner(session_id) - if not await self.has_session_async(session_id): + try: + session = await self.get_session_async(session_id) + except KeyError: return ApprovalResponse( session_id=session_id, request_id=approval_req.request_id, approved=False, feedback="Session not found", ) - - session = await self.get_session_async(session_id) + await self._assert_owner(session_id) if not session.turn_in_progress: return ApprovalResponse( session_id=session_id, request_id=approval_req.request_id, approved=False,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@coding-agent/src/coding_agent/ui/session_manager.py` around lines 1166 - 1175, The current logic calls _assert_owner(session_id) before checking existence, causing a RuntimeError when a session has no owner and making the "Session not found" ApprovalResponse unreachable; swap the checks so you call has_session_async(session_id) first and return the ApprovalResponse(...) if false, and only then call _assert_owner(session_id) and proceed to get_session_async(session_id) — update the block around _assert_owner, has_session_async, get_session_async and the ApprovalResponse/approval_req usage accordingly.
1271-1307:⚠️ Potential issue | 🟠 Major
ensure_session_runtime()is still only point-in-time fenced.This method validates ownership once, then awaits
_restore_tape()andadapter.initialize()before writingruntime_pipeline,runtime_ctx,runtime_adapter, andtape_id. If the lease changes hands during either await, the stale owner still installs a runtime and persists it. The other destructive entry points were moved behind local locking; this path needs the same revalidation/serialization around the commit step.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@coding-agent/src/coding_agent/ui/session_manager.py`:
- Around line 969-970: The new RuntimeError from _assert_owner should not
surface as a 500 in the HTTP layer; in close_session() wrap the await
self._assert_owner(session_id) call in a try/except that catches RuntimeError
and re-raises a KeyError (or another exception already mapped to a
client-visible conflict) with the original message (e.g. "raise KeyError(str(e))
from e") so the existing http_server handling will map the stale-owner/lease
conflict to the proper client response instead of an internal server error.
---
Outside diff comments:
In `@coding-agent/src/coding_agent/ui/session_manager.py`:
- Around line 1166-1175: The current logic calls _assert_owner(session_id)
before checking existence, causing a RuntimeError when a session has no owner
and making the "Session not found" ApprovalResponse unreachable; swap the checks
so you call has_session_async(session_id) first and return the
ApprovalResponse(...) if false, and only then call _assert_owner(session_id) and
proceed to get_session_async(session_id) — update the block around
_assert_owner, has_session_async, get_session_async and the
ApprovalResponse/approval_req usage accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a3ee12e6-9c4b-474f-ba8a-8de518247c31
📒 Files selected for processing (2)
coding-agent/src/coding_agent/ui/session_manager.pycoding-agent/tests/ui/test_session_manager_owner_checks.py
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This was referenced Apr 21, 2026
This was referenced Jun 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Verification
uv run pytest tests/ui/test_session_owner_store.py tests/ui/test_session_manager_owner_checks.py -vuv run pytest tests/agentkit/storage/test_pg.py -k "owner or lease" -vSummary by CodeRabbit
New Features
Tests