wehub-resource-sync
824 lines
37 KiB
Python
824 lines
37 KiB
Python
"""Tests for the ``WS /v1/sessions/updates`` push stream.
|
|
|
|
The endpoint replaces the web app's HTTP poll of ``GET /v1/sessions``:
|
|
a client sends the session ids it's watching and the server pushes a
|
|
``snapshot`` followed by ``changed`` / ``removed`` deltas as those
|
|
sessions change. These tests drive the real route (no mocks of the
|
|
store or auth) against file-backed SQLite stores, mutating persisted
|
|
conversation state to trigger deltas, and assert cross-user isolation.
|
|
|
|
The per-connection rescan interval is monkeypatched down so the
|
|
interval-driven ``changed`` / ``removed`` frames arrive promptly
|
|
instead of after the production 4 s cadence.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
import omnigent.server.routes.sessions as sessions_routes
|
|
from omnigent.server.auth import LEVEL_OWNER, UnifiedAuthProvider
|
|
from omnigent.server.routes.sessions import SessionLiveness, create_sessions_router
|
|
from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore
|
|
from omnigent.stores.comment_store.sqlalchemy_store import SqlAlchemyCommentStore
|
|
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
|
|
from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore
|
|
|
|
ALICE = "alice@example.com"
|
|
BOB = "bob@example.com"
|
|
|
|
|
|
class _NoIdentityAuthProvider:
|
|
"""Auth provider whose handshake yields no identity.
|
|
|
|
Exercises the updates-stream's reject-when-unauthenticated gate
|
|
deterministically — unlike header mode, which falls back to the
|
|
reserved ``"local"`` user and would never return ``None``.
|
|
"""
|
|
|
|
def get_user_id(self, request: object) -> None:
|
|
"""Always return ``None`` (no authenticated identity)."""
|
|
del request
|
|
return
|
|
|
|
|
|
@pytest.fixture
|
|
def fast_rescan(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
Shrink the per-connection rescan interval so interval-driven
|
|
deltas (``changed`` / ``removed``) land in well under a second.
|
|
|
|
Patches the module-level constants the ticker reads at sleep time;
|
|
these are plain config values, not ``asyncio`` singletons, so the
|
|
patch is local and safe. The heartbeat interval is also shrunk so
|
|
idle-only tests receive heartbeats promptly instead of waiting the
|
|
production 30 s.
|
|
"""
|
|
monkeypatch.setattr(sessions_routes, "_SESSION_UPDATES_RESCAN_INTERVAL_S", 0.05)
|
|
monkeypatch.setattr(sessions_routes, "_SESSION_UPDATES_HEARTBEAT_INTERVAL_S", 0.1)
|
|
|
|
|
|
@pytest.fixture
|
|
def stores(
|
|
db_uri: str,
|
|
) -> tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore]:
|
|
"""Real file-backed stores so writes from the test thread are visible
|
|
to the WS handler thread."""
|
|
return (
|
|
SqlAlchemyConversationStore(db_uri),
|
|
SqlAlchemyAgentStore(db_uri),
|
|
SqlAlchemyPermissionStore(db_uri),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def liveness_state() -> dict[str, SessionLiveness]:
|
|
"""Mutable liveness map the test can flip mid-connection to drive a
|
|
``runner_online`` / ``host_online`` change through the stream. Ids
|
|
absent from the map default to ``runner_online=True``,
|
|
``host_online=None`` (matching the server's missing-row terminal)."""
|
|
return {}
|
|
|
|
|
|
@pytest.fixture
|
|
def comment_store(db_uri: str) -> SqlAlchemyCommentStore:
|
|
"""Real file-backed comment store so comment writes from the test
|
|
thread are visible to the WS handler thread, mirroring ``stores``."""
|
|
return SqlAlchemyCommentStore(db_uri)
|
|
|
|
|
|
@pytest.fixture
|
|
def app(
|
|
stores: tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore],
|
|
liveness_state: dict[str, SessionLiveness],
|
|
comment_store: SqlAlchemyCommentStore,
|
|
) -> FastAPI:
|
|
"""Minimal app mounting only the sessions router, with header-based
|
|
auth and a real permission store — the surface the updates stream
|
|
actually exercises. ``liveness_lookup`` reads the mutable
|
|
``liveness_state`` (default runner online, no host) so tests control
|
|
liveness."""
|
|
conversation_store, agent_store, permission_store = stores
|
|
|
|
def _liveness_lookup(ids: list[str]) -> dict[str, SessionLiveness]:
|
|
return {
|
|
sid: liveness_state.get(sid, SessionLiveness(runner_online=True, host_online=None))
|
|
for sid in ids
|
|
}
|
|
|
|
app = FastAPI()
|
|
app.include_router(
|
|
create_sessions_router(
|
|
conversation_store=conversation_store,
|
|
agent_store=agent_store,
|
|
auth_provider=UnifiedAuthProvider(source="header"),
|
|
permission_store=permission_store,
|
|
liveness_lookup=_liveness_lookup,
|
|
comment_store=comment_store,
|
|
),
|
|
prefix="/v1",
|
|
)
|
|
return app
|
|
|
|
|
|
def _seed_session(
|
|
stores: tuple[SqlAlchemyConversationStore, SqlAlchemyAgentStore, SqlAlchemyPermissionStore],
|
|
*,
|
|
owner: str,
|
|
title: str,
|
|
) -> str:
|
|
"""Create a session-shaped conversation (non-null ``agent_id``) owned
|
|
by ``owner`` and return its id."""
|
|
conversation_store, agent_store, permission_store = stores
|
|
# The conversations.agent_id FK requires a real agent row; create one
|
|
# idempotently (the same agent backs every seeded session).
|
|
if agent_store.get("ag_test") is None:
|
|
agent_store.create(
|
|
agent_id="ag_test",
|
|
name="test-agent",
|
|
bundle_location="ag_test/bundle",
|
|
)
|
|
conv = conversation_store.create_conversation(title=title, agent_id="ag_test")
|
|
permission_store.ensure_user(owner)
|
|
permission_store.grant(owner, conv.id, LEVEL_OWNER)
|
|
return conv.id
|
|
|
|
|
|
def _recv_until(ws: object, wanted: set[str], *, max_frames: int = 50) -> dict[str, object]:
|
|
"""
|
|
Read frames until one whose ``type`` is in ``wanted`` arrives.
|
|
|
|
Heartbeats and snapshots that aren't being awaited are skipped so a
|
|
test can target a specific delta frame without depending on cadence.
|
|
|
|
:param ws: The connected test WebSocket.
|
|
:param wanted: Frame ``type`` values to stop on, e.g. ``{"changed"}``.
|
|
:param max_frames: Safety bound so a missing frame fails fast rather
|
|
than hanging.
|
|
:returns: The first matching frame as a parsed dict.
|
|
"""
|
|
for _ in range(max_frames):
|
|
frame = json.loads(ws.receive_text()) # type: ignore[attr-defined]
|
|
if frame.get("type") in wanted:
|
|
return frame
|
|
raise AssertionError(f"no frame in {wanted} after {max_frames} frames")
|
|
|
|
|
|
def test_watch_returns_snapshot_of_accessible_sessions(app: FastAPI, stores) -> None:
|
|
"""A ``watch`` for owned ids returns a snapshot containing exactly
|
|
those sessions, with their persisted titles."""
|
|
s1 = _seed_session(stores, owner=ALICE, title="first")
|
|
s2 = _seed_session(stores, owner=ALICE, title="second")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1, s2]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
items = {item["id"]: item for item in snapshot["items"]} # type: ignore[index]
|
|
# Both watched, owned sessions are present — proves the snapshot
|
|
# reads real conversation rows for exactly the watched ids.
|
|
assert set(items) == {s1, s2}
|
|
assert items[s1]["title"] == "first"
|
|
assert items[s2]["title"] == "second"
|
|
# No relay has reported, so the relay-fed status cache misses and
|
|
# the list status presents as idle.
|
|
assert items[s1]["status"] == "idle"
|
|
# Both liveness fields are folded into the payload (defaults:
|
|
# runner online, no host) so the client can drop its /health poll
|
|
# for these sessions. A regression that stopped applying host_online
|
|
# would drop the key (stream dumps full rows, so it must be present).
|
|
assert items[s1]["runner_online"] is True
|
|
assert items[s1]["host_online"] is None
|
|
|
|
|
|
def test_child_busy_rollup_flows_through_updates_stream(
|
|
app: FastAPI,
|
|
stores,
|
|
fast_rescan: None,
|
|
) -> None:
|
|
"""
|
|
A watched parent row reflects direct child sub-agent busy status.
|
|
|
|
The sidebar relies on this WebSocket stream after the initial list
|
|
fetch, so the stream's shared list-item builder must pass the same
|
|
direct child ids as ``GET /v1/sessions``. This seeds a real child
|
|
conversation and the relay-fed status cache, then proves both the
|
|
initial snapshot and a later diff frame carry the rolled-up parent
|
|
status.
|
|
"""
|
|
conversation_store = stores[0]
|
|
parent_id = _seed_session(stores, owner=ALICE, title="parent")
|
|
child = conversation_store.create_conversation(
|
|
kind="sub_agent",
|
|
title="coder:auth",
|
|
parent_conversation_id=parent_id,
|
|
agent_id="ag_test",
|
|
)
|
|
sessions_routes._session_status_cache.pop(parent_id, None)
|
|
sessions_routes._session_status_cache[child.id] = "waiting"
|
|
try:
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [parent_id]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
snapshot_items = {item["id"]: item for item in snapshot["items"]} # type: ignore[index]
|
|
assert snapshot_items[parent_id]["status"] == "running"
|
|
|
|
sessions_routes._session_status_cache[child.id] = "idle"
|
|
changed = _recv_until(ws, {"changed"})
|
|
changed_items = {item["id"]: item for item in changed["items"]} # type: ignore[index]
|
|
assert changed_items[parent_id]["status"] == "idle"
|
|
finally:
|
|
sessions_routes._session_status_cache.pop(parent_id, None)
|
|
sessions_routes._session_status_cache.pop(child.id, None)
|
|
|
|
|
|
def test_runner_offline_pushes_changed_frame(
|
|
app: FastAPI, stores, liveness_state: dict[str, SessionLiveness], fast_rescan: None
|
|
) -> None:
|
|
"""Flipping a watched session's runner to offline (host still up) pushes
|
|
a ``changed`` frame carrying ``runner_online: false`` and
|
|
``host_online: true`` — this is the runner-down-but-host-alive state the
|
|
open view uses, and it lets the client retire its /health poll."""
|
|
s1 = _seed_session(stores, owner=ALICE, title="live")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
assert snapshot["items"][0]["runner_online"] is True # type: ignore[index]
|
|
# Runner goes offline but the host stays alive; the next rescan must
|
|
# diff both fields and push them. A broken fold (a field missing or
|
|
# not diffed) would yield only heartbeats and hang this wait.
|
|
liveness_state[s1] = SessionLiveness(runner_online=False, host_online=True)
|
|
changed = _recv_until(ws, {"changed"})
|
|
items = {item["id"]: item for item in changed["items"]} # type: ignore[index]
|
|
# Strict runner_online flipped to False while host_online reports the
|
|
# host is still reachable — exactly the pair that distinguishes "wake
|
|
# the runner with a message" from "host offline, must reconnect".
|
|
assert items[s1]["runner_online"] is False
|
|
assert items[s1]["host_online"] is True
|
|
|
|
|
|
def test_list_sessions_omits_liveness_fields(
|
|
app: FastAPI, stores, liveness_state: dict[str, SessionLiveness]
|
|
) -> None:
|
|
"""``GET /v1/sessions`` does NOT compute per-item liveness.
|
|
|
|
The list deliberately skips ``runner_online`` / ``host_online``: no
|
|
list consumer reads them (the sidebar no longer surfaces connection
|
|
state), and the open-session view sources liveness from the
|
|
single-session snapshot, the WS updates stream, and the ``/health``
|
|
poll. Skipping it here drops the session-connectivity and hosts-table
|
|
queries from every list call. The updates stream still carries the
|
|
fields (see the changed-frame tests). Even with a liveness state set
|
|
for the session, the list must not surface it."""
|
|
s1 = _seed_session(stores, owner=ALICE, title="live")
|
|
liveness_state[s1] = SessionLiveness(runner_online=False, host_online=True)
|
|
|
|
resp = TestClient(app).get("/v1/sessions", headers={"X-Forwarded-Email": ALICE})
|
|
|
|
assert resp.status_code == 200
|
|
items = {item["id"]: item for item in resp.json()["data"]}
|
|
# The list uses exclude_none; both keys are absent because the list
|
|
# never computes liveness. Re-adding the lookup here would reintroduce
|
|
# the per-list connectivity + hosts queries this change removed.
|
|
assert "runner_online" not in items[s1]
|
|
assert "host_online" not in items[s1]
|
|
|
|
|
|
def test_title_change_pushes_changed_frame(app: FastAPI, stores, fast_rescan: None) -> None:
|
|
"""Mutating a watched session's persisted title makes the server push
|
|
a ``changed`` frame carrying the new value."""
|
|
conversation_store = stores[0]
|
|
s1 = _seed_session(stores, owner=ALICE, title="before")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# Mutate persisted state from the test thread; the handler's next
|
|
# rescan must observe it and diff it against the snapshot baseline.
|
|
conversation_store.update_conversation(s1, title="after")
|
|
changed = _recv_until(ws, {"changed"})
|
|
items = {item["id"]: item for item in changed["items"]} # type: ignore[index]
|
|
# The changed frame carries the new title — proves the diff fired
|
|
# on the real field change, not a spurious or stale value.
|
|
assert items[s1]["title"] == "after"
|
|
|
|
|
|
def test_cleared_field_pushes_explicit_null(app: FastAPI, stores, fast_rescan: None) -> None:
|
|
"""Clearing a previously-set nullable field (a runner unbind nulling
|
|
``runner_id``) pushes a ``changed`` frame that carries the key as an
|
|
explicit ``null``.
|
|
|
|
The stream dumps full rows (NOT ``exclude_none``) precisely so a
|
|
non-null → null transition arrives as an explicit ``null`` the client can
|
|
overlay-clear, rather than a dropped key that would leave the stale value.
|
|
If the stream regressed to ``exclude_none``, ``runner_id`` would be absent
|
|
from the frame and this test's ``in`` assertion fails.
|
|
"""
|
|
conversation_store = stores[0]
|
|
s1 = _seed_session(stores, owner=ALICE, title="bound")
|
|
# Bind a runner first so the snapshot baseline carries runner_id; the
|
|
# clear below is then a genuine non-null → null transition.
|
|
assert conversation_store.set_runner_id(s1, "rnr_test") is True
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
# Baseline carries the bound runner — the value the clear must reverse.
|
|
assert snapshot["items"][0]["runner_id"] == "rnr_test" # type: ignore[index]
|
|
conversation_store.clear_runner_id(s1)
|
|
changed = _recv_until(ws, {"changed"})
|
|
item = {i["id"]: i for i in changed["items"]}[s1] # type: ignore[index]
|
|
# The key is present AND explicitly null. Presence proves the stream
|
|
# dumps full rows (exclude_none would omit it); the None value proves
|
|
# the client can overlay-clear the stale runner_id.
|
|
assert "runner_id" in item
|
|
assert item["runner_id"] is None
|
|
|
|
|
|
def test_no_change_emits_no_changed_frame(app: FastAPI, stores, fast_rescan: None) -> None:
|
|
"""An idle watched session produces no ``changed`` frames — only
|
|
heartbeats — so a static list generates no diff traffic."""
|
|
s1 = _seed_session(stores, owner=ALICE, title="static")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# Several rescans elapse (interval 0.05 s); none should diff since
|
|
# nothing changed. Any "changed"/"removed" here means the differ
|
|
# is emitting on unchanged state.
|
|
for _ in range(5):
|
|
frame = json.loads(ws.receive_text())
|
|
assert frame["type"] == "heartbeat", (
|
|
f"expected only heartbeats for an unchanged session, got {frame['type']}"
|
|
)
|
|
|
|
|
|
def test_delete_pushes_removed_frame(app: FastAPI, stores, fast_rescan: None) -> None:
|
|
"""Deleting a watched session makes the server push a ``removed``
|
|
frame for its id (the row no longer resolves)."""
|
|
permission_store = stores[2]
|
|
s1 = _seed_session(stores, owner=ALICE, title="doomed")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# Drop access so the next rescan can no longer resolve the id for
|
|
# this user — the list-stream's definition of "removed".
|
|
permission_store.revoke(ALICE, s1)
|
|
removed = _recv_until(ws, {"removed"})
|
|
# The removed frame names exactly the now-inaccessible id.
|
|
assert removed["ids"] == [s1]
|
|
|
|
|
|
def test_other_users_session_is_not_visible(app: FastAPI, stores, fast_rescan: None) -> None:
|
|
"""Bob watching Alice's session never receives it — neither in the
|
|
snapshot nor in any later frame (cross-user isolation, W-series)."""
|
|
s_alice = _seed_session(stores, owner=ALICE, title="alice-only")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": BOB}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s_alice]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
# Bob has no grant on Alice's session → the snapshot omits it.
|
|
assert snapshot["items"] == []
|
|
# And it must not leak via a later changed frame when Alice mutates
|
|
# her own session.
|
|
stores[0].update_conversation(s_alice, title="changed-by-alice")
|
|
for _ in range(5):
|
|
frame = json.loads(ws.receive_text())
|
|
assert frame["type"] == "heartbeat", (
|
|
f"Bob received a {frame['type']} frame for a session he "
|
|
f"cannot access — cross-user leak"
|
|
)
|
|
|
|
|
|
def test_unauthenticated_connection_is_rejected(stores) -> None:
|
|
"""With permissions enabled, a socket whose handshake yields no
|
|
identity is closed at the handshake (policy violation) before any
|
|
session data is read."""
|
|
conversation_store, agent_store, permission_store = stores
|
|
app = FastAPI()
|
|
app.include_router(
|
|
create_sessions_router(
|
|
conversation_store=conversation_store,
|
|
agent_store=agent_store,
|
|
auth_provider=_NoIdentityAuthProvider(), # type: ignore[arg-type]
|
|
permission_store=permission_store,
|
|
),
|
|
prefix="/v1",
|
|
)
|
|
with pytest.raises(WebSocketDisconnect) as exc_info:
|
|
with TestClient(app).websocket_connect("/v1/sessions/updates"):
|
|
pass
|
|
# 1008 = WS_1008_POLICY_VIOLATION — the auth gate fired before accept.
|
|
assert exc_info.value.code == 1008
|
|
|
|
|
|
def test_watch_set_truncated_at_cap(
|
|
app: FastAPI,
|
|
stores,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
caplog: pytest.LogCaptureFixture,
|
|
) -> None:
|
|
"""A watch-set larger than the cap is truncated to the cap, and the
|
|
drop is logged. Ids past the cap never appear in the snapshot — they
|
|
fall back to the client's list poll, so the server must surface the
|
|
truncation rather than silently shrinking the set."""
|
|
# Shrink the cap to 2 so three seeded sessions exceed it without
|
|
# seeding the production 500.
|
|
monkeypatch.setattr(sessions_routes, "_SESSION_UPDATES_MAX_WATCHED", 2)
|
|
s1 = _seed_session(stores, owner=ALICE, title="one")
|
|
s2 = _seed_session(stores, owner=ALICE, title="two")
|
|
s3 = _seed_session(stores, owner=ALICE, title="three")
|
|
with caplog.at_level(logging.WARNING, logger="omnigent.server.routes.sessions"):
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
# Send all three in order; the reader keeps the first two
|
|
# (preserving order) and drops the third.
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1, s2, s3]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
ids = {item["id"] for item in snapshot["items"]} # type: ignore[index]
|
|
# Exactly the first two watched ids resolve; the third was
|
|
# truncated before the snapshot fetch ever saw it. A regression
|
|
# that dropped the cap would let all three through here.
|
|
assert ids == {s1, s2}
|
|
assert len(snapshot["items"]) == 2 # type: ignore[arg-type]
|
|
# The truncation is logged with the distinct-id count so an oversized
|
|
# watch-set is diagnosable. Absence here means the drop went silent.
|
|
assert any("watch-set truncated" in record.getMessage() for record in caplog.records), (
|
|
"expected a watch-set-truncation warning, got none"
|
|
)
|
|
|
|
|
|
def test_transient_store_error_does_not_kill_stream(
|
|
app: FastAPI,
|
|
stores,
|
|
fast_rescan: None,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A store read that raises during one rescan tick is logged and
|
|
skipped — the stream stays open and a later real change still pushes a
|
|
``changed`` frame. Without the ticker's guard, the raised exception
|
|
would tear the connection down and force a reconnect + re-snapshot."""
|
|
conversation_store = stores[0]
|
|
s1 = _seed_session(stores, owner=ALICE, title="before")
|
|
real_get = conversation_store.get_conversations
|
|
calls = {"n": 0}
|
|
|
|
def flaky_get(ids: list[str]) -> dict[str, object]:
|
|
"""Succeed on the connect snapshot, fail the first rescan tick,
|
|
then recover — a transient blip, not a permanent outage."""
|
|
calls["n"] += 1
|
|
# Call 1 is the snapshot (must succeed so the client gets a
|
|
# baseline). Call 2 is the first ticker rescan — fail it.
|
|
if calls["n"] == 2:
|
|
raise RuntimeError("transient store blip")
|
|
return real_get(ids)
|
|
|
|
monkeypatch.setattr(conversation_store, "get_conversations", flaky_get)
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# Mutate after the failing tick is guaranteed to occur (call 2);
|
|
# a later successful tick must observe it. If the failed tick had
|
|
# killed the ticker, the connection would close and _recv_until
|
|
# would raise WebSocketDisconnect instead of returning a frame.
|
|
conversation_store.update_conversation(s1, title="after")
|
|
changed = _recv_until(ws, {"changed"})
|
|
items = {item["id"]: item for item in changed["items"]} # type: ignore[index]
|
|
assert items[s1]["title"] == "after"
|
|
# >= 3 proves at least one rescan ran after the injected failure
|
|
# (call 1 = snapshot, call 2 = failed tick, call 3+ = recovered ticks).
|
|
assert calls["n"] >= 3, (
|
|
f"expected the ticker to keep rescanning after a failed tick, "
|
|
f"but get_conversations was called only {calls['n']} time(s)"
|
|
)
|
|
|
|
|
|
def test_session_added_event_pushes_unwatched_session(
|
|
app: FastAPI, stores, fast_rescan: None
|
|
) -> None:
|
|
"""A ``session_added`` discovery event pushes a session the client isn't
|
|
watching — the path that lets a session created elsewhere (another tab /
|
|
CLI) enter the sidebar without a list poll.
|
|
|
|
The client can never put an unknown session in its watch-set, so the
|
|
interval diff can't surface it; the server reacts to the create/grant
|
|
event instead. The pushed frame carries the real persisted row (fetched,
|
|
not echoed from the event), which is what the client reconciles into the
|
|
list. ``fast_rescan`` shrinks the heartbeat cadence so that if the push
|
|
regressed, ``_recv_until`` exhausts its frame budget on heartbeats and
|
|
fails fast instead of blocking on the 30 s production heartbeat.
|
|
"""
|
|
s1 = _seed_session(stores, owner=ALICE, title="watched")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# A brand-new session is created for Alice elsewhere and announced on
|
|
# her discovery channel — it is NOT in the watch-set above.
|
|
s2 = _seed_session(stores, owner=ALICE, title="brand new")
|
|
sessions_routes.user_session_stream.publish(
|
|
ALICE, {"type": "session_added", "session_id": s2}
|
|
)
|
|
changed = _recv_until(ws, {"changed"})
|
|
items = {item["id"]: item for item in changed["items"]} # type: ignore[index]
|
|
# The unwatched new session is pushed, carrying its persisted title —
|
|
# proving the server fetched the real row, not echoed the event.
|
|
assert s2 in items
|
|
assert items[s2]["title"] == "brand new"
|
|
|
|
|
|
def test_session_added_for_inaccessible_session_is_not_pushed(
|
|
app: FastAPI, stores, fast_rescan: None
|
|
) -> None:
|
|
"""A discovery announcement for a session the user can't access is dropped.
|
|
|
|
The per-id access check in the push path is the cross-user safety net: even
|
|
if a ``session_added`` for Bob's session landed on Alice's channel (e.g. a
|
|
future mis-keyed publish), Alice must never be pushed a session she has no
|
|
grant on. A regression that skipped the access check would leak Bob's row
|
|
into Alice's stream as a ``changed`` frame.
|
|
"""
|
|
s1 = _seed_session(stores, owner=ALICE, title="alice")
|
|
bob_session = _seed_session(stores, owner=BOB, title="bob only")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
_recv_until(ws, {"snapshot"})
|
|
# Announce Bob's session on Alice's channel; the access check must drop it.
|
|
sessions_routes.user_session_stream.publish(
|
|
ALICE, {"type": "session_added", "session_id": bob_session}
|
|
)
|
|
# Alice's watched session is static, so every frame must be a heartbeat;
|
|
# a ``changed`` here (carrying bob_session) would be a cross-user leak.
|
|
for _ in range(5):
|
|
frame = json.loads(ws.receive_text())
|
|
assert frame["type"] == "heartbeat", (
|
|
f"expected only heartbeats, got {frame['type']} "
|
|
f"(a changed frame would mean Bob's session leaked to Alice)"
|
|
)
|
|
|
|
|
|
# ── comments fingerprint ──────────────────────────────────────────────
|
|
|
|
|
|
# One second in microseconds — comments_updated_at is epoch-µs.
|
|
_US = 1_000_000
|
|
|
|
|
|
@pytest.fixture
|
|
def comment_clock(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]:
|
|
"""Deterministic write clock for comment rows.
|
|
|
|
Patches the ``now_epoch_us`` reference the comment store uses for
|
|
all comment timestamps (``created_at`` derives from the same read)
|
|
so consecutive comment writes get distinct, controlled values.
|
|
|
|
:param monkeypatch: Pytest monkeypatch fixture.
|
|
:returns: Mutable ``{"now": <epoch seconds>}`` dict; assign
|
|
``clock["now"]`` to move time forward.
|
|
"""
|
|
state = {"now": 1_000}
|
|
monkeypatch.setattr(
|
|
"omnigent.stores.comment_store.sqlalchemy_store.now_epoch_us",
|
|
lambda: state["now"] * _US,
|
|
)
|
|
return state
|
|
|
|
|
|
def test_comment_add_pushes_changed_frame(
|
|
app: FastAPI,
|
|
stores,
|
|
comment_store: SqlAlchemyCommentStore,
|
|
comment_clock: dict[str, int],
|
|
fast_rescan: None,
|
|
) -> None:
|
|
"""Adding a comment to a watched session pushes a ``changed`` frame.
|
|
|
|
This is the freshness path the CommentsPanel relies on: another
|
|
user's POST /comments must surface on this stream so the web app
|
|
can invalidate its cached comment list without polling.
|
|
"""
|
|
s1 = _seed_session(stores, owner=ALICE, title="commented")
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
items = {item["id"]: item for item in snapshot["items"]}
|
|
# Baseline: no comments yet. The stream dumps full rows, so the
|
|
# fingerprint fields must be present (count 0 / explicit null).
|
|
assert items[s1]["comments_count"] == 0
|
|
assert items[s1]["comments_updated_at"] is None
|
|
|
|
comment_store.add(
|
|
conversation_id=s1,
|
|
path="src/app.py",
|
|
body="needs a null check",
|
|
start_index=0,
|
|
end_index=10,
|
|
)
|
|
changed = _recv_until(ws, {"changed"})
|
|
changed_items = {item["id"]: item for item in changed["items"]}
|
|
# The add must move both fingerprint fields; a frame without them
|
|
# means the WS builder isn't reading the comment store.
|
|
assert changed_items[s1]["comments_count"] == 1
|
|
assert changed_items[s1]["comments_updated_at"] == 1_000 * _US
|
|
|
|
|
|
def test_comment_status_change_pushes_changed_frame(
|
|
app: FastAPI,
|
|
stores,
|
|
comment_store: SqlAlchemyCommentStore,
|
|
comment_clock: dict[str, int],
|
|
fast_rescan: None,
|
|
) -> None:
|
|
"""Marking a comment addressed pushes a ``changed`` frame.
|
|
|
|
This is the agent flow: ``update_comment`` (status draft →
|
|
addressed) changes no row count, so the ``updated_at`` bump is the
|
|
only signal — if it didn't move, other viewers would keep showing
|
|
the comment as open.
|
|
"""
|
|
s1 = _seed_session(stores, owner=ALICE, title="addressed")
|
|
comment = comment_store.add(
|
|
conversation_id=s1,
|
|
path="src/app.py",
|
|
body="fix me",
|
|
start_index=0,
|
|
end_index=6,
|
|
)
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
items = {item["id"]: item for item in snapshot["items"]}
|
|
assert items[s1]["comments_count"] == 1
|
|
assert items[s1]["comments_updated_at"] == 1_000 * _US
|
|
|
|
comment_clock["now"] = 2_000
|
|
comment_store.update_comment(comment.id, s1, status="addressed")
|
|
changed = _recv_until(ws, {"changed"})
|
|
changed_items = {item["id"]: item for item in changed["items"]}
|
|
# Count is unchanged — the moved timestamp alone must trigger the
|
|
# frame. A stale creation-time value here means update_comment didn't bump
|
|
# updated_at and in-place edits are invisible to clients.
|
|
assert changed_items[s1]["comments_count"] == 1
|
|
assert changed_items[s1]["comments_updated_at"] == 2_000 * _US
|
|
|
|
|
|
def test_comment_delete_of_older_comment_pushes_changed_frame(
|
|
app: FastAPI,
|
|
stores,
|
|
comment_store: SqlAlchemyCommentStore,
|
|
comment_clock: dict[str, int],
|
|
fast_rescan: None,
|
|
) -> None:
|
|
"""Deleting a non-newest comment pushes a ``changed`` frame.
|
|
|
|
The deleted row is not the most recently updated one, so
|
|
``max(updated_at)`` is unchanged — the count drop is the only
|
|
signal. This is the case that justifies ``comments_count``.
|
|
"""
|
|
s1 = _seed_session(stores, owner=ALICE, title="deleted")
|
|
older = comment_store.add(
|
|
conversation_id=s1, path="a.py", body="old", start_index=0, end_index=3
|
|
)
|
|
comment_clock["now"] = 2_000
|
|
comment_store.add(conversation_id=s1, path="a.py", body="new", start_index=4, end_index=7)
|
|
with TestClient(app).websocket_connect(
|
|
"/v1/sessions/updates", headers={"X-Forwarded-Email": ALICE}
|
|
) as ws:
|
|
ws.send_text(json.dumps({"type": "watch", "session_ids": [s1]}))
|
|
snapshot = _recv_until(ws, {"snapshot"})
|
|
items = {item["id"]: item for item in snapshot["items"]}
|
|
assert items[s1]["comments_count"] == 2
|
|
assert items[s1]["comments_updated_at"] == 2_000 * _US
|
|
|
|
comment_store.delete(older.id, s1)
|
|
changed = _recv_until(ws, {"changed"})
|
|
changed_items = {item["id"]: item for item in changed["items"]}
|
|
# The timestamp must NOT move (the surviving comment is the
|
|
# newest) — count 2 → 1 is what fires the frame. If this hangs
|
|
# at _recv_until, deletes of older comments are invisible.
|
|
assert changed_items[s1]["comments_count"] == 1
|
|
assert changed_items[s1]["comments_updated_at"] == 2_000 * _US
|
|
|
|
|
|
def test_daily_cost_recorded_for_owned_session_without_a_policy(app: FastAPI, stores) -> None:
|
|
"""Per-user daily cost is recorded even when the session has no policy.
|
|
|
|
The policy gate that previously suppressed the ``user_daily_cost`` write on
|
|
no-policy sessions has been removed, so the daily rollup accrues for any
|
|
owned session that records spend. Here ``ag_test`` has no guardrails/policy,
|
|
yet posting cumulative spend must still land in the owner's daily rollup. A
|
|
regression that re-added the gate would leave ``get_daily_cost`` at 0.0.
|
|
"""
|
|
conversation_store = stores[0]
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
sid = _seed_session(stores, owner=ALICE, title="no-policy session")
|
|
resp = TestClient(app).post(
|
|
f"/v1/sessions/{sid}/events",
|
|
json={"type": "external_session_usage", "data": {"cumulative_cost_usd": 0.5}},
|
|
headers={"X-Forwarded-Email": ALICE},
|
|
)
|
|
assert resp.status_code == 202, resp.text
|
|
# Recorded under the owner despite no policy on the session — proves the
|
|
# write is no longer policy-gated.
|
|
assert conversation_store.get_daily_cost(ALICE, today) == pytest.approx(0.5)
|
|
|
|
|
|
def test_daily_cost_tracks_display_cost_not_policy_cost(app: FastAPI, stores) -> None:
|
|
"""The daily rollup follows ``total_cost_usd`` (S), not ``policy_cost_usd``.
|
|
|
|
claude-native posts the statusLine total ``S`` as ``cumulative_cost_usd``
|
|
(display) and a higher ``max(S, C)`` as ``policy_cost_usd`` (the real-time
|
|
gate value, inflated mid-turn by in-flight sub-agent spend). The per-user
|
|
daily report must reflect real spend, so its delta comes from
|
|
``total_cost_usd`` = ``S``. A regression that fed ``policy_cost_usd`` into
|
|
the rollup would over-report the daily total by the gate's mid-turn lead.
|
|
"""
|
|
conversation_store = stores[0]
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
sid = _seed_session(stores, owner=ALICE, title="cost split session")
|
|
resp = TestClient(app).post(
|
|
f"/v1/sessions/{sid}/events",
|
|
json={
|
|
"type": "external_session_usage",
|
|
# Display S = 0.20; enforcement value 0.90 leads it (sub-agent mid-run).
|
|
"data": {"cumulative_cost_usd": 0.20, "policy_cost_usd": 0.90},
|
|
},
|
|
headers={"X-Forwarded-Email": ALICE},
|
|
)
|
|
assert resp.status_code == 202, resp.text
|
|
# Daily = S (0.20), NOT the 0.90 enforcement figure. A value of 0.90 would
|
|
# mean the rollup tracked policy_cost_usd and inherited the gate's
|
|
# mid-turn inflation — the daily over-report this split prevents.
|
|
assert conversation_store.get_daily_cost(ALICE, today) == pytest.approx(0.20)
|
|
|
|
|
|
def test_daily_cost_attributed_via_root_for_sub_agent_without_owner_grant(
|
|
app: FastAPI, stores
|
|
) -> None:
|
|
"""Sub-agent spend is attributed to the root session's owner.
|
|
|
|
Relay / SDK sub-agents are spawned by the internal runner (no user
|
|
context in the POST), so their conversations never receive an owner
|
|
permission grant. Previously ``_record_daily_cost`` called
|
|
``get_session_owner(child.id)``, got ``None``, and silently dropped the
|
|
cost from the daily rollup — the per-user daily budget never saw it.
|
|
|
|
The fix: fall back to ``get_session_owner(root_conversation_id)`` when
|
|
the direct lookup misses. This test creates a parent (owned) + a child
|
|
conversation (no grant, but ``root_conversation_id`` → parent), posts
|
|
cumulative spend on the child, and asserts the owner's daily total rises.
|
|
"""
|
|
conversation_store, _agent_store, _permission_store = stores
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
|
|
# Parent session — owned by Alice.
|
|
parent_id = _seed_session(stores, owner=ALICE, title="parent session")
|
|
|
|
# Child conversation — simulates a relay sub-agent: no permission grant.
|
|
# Passing parent_conversation_id causes create_conversation to inherit
|
|
# root_conversation_id from the parent automatically.
|
|
child = conversation_store.create_conversation(
|
|
title="sub-agent",
|
|
agent_id="ag_test",
|
|
parent_conversation_id=parent_id,
|
|
)
|
|
# Sanity: child has no owner grant (the gap being fixed).
|
|
assert conversation_store.get_session_owner(child.id) is None
|
|
|
|
# Post cumulative cost on the child — no auth header (internal runner path).
|
|
resp = TestClient(app).post(
|
|
f"/v1/sessions/{child.id}/events",
|
|
json={"type": "external_session_usage", "data": {"cumulative_cost_usd": 0.75}},
|
|
)
|
|
assert resp.status_code == 202, resp.text
|
|
|
|
# Sub-agent spend must appear in Alice's daily rollup via the root fallback.
|
|
assert conversation_store.get_daily_cost(ALICE, today) == pytest.approx(0.75)
|