fix(openviking): implement on_session_switch hook (#28296)
OpenVikingMemoryProvider only overrides on_session_end and inherits the base-class no-op for on_session_switch. When the agent rotates session_id (via /new, /branch, /reset, /resume, or context compression), the provider's cached _session_id stays at the value initialize() captured. All subsequent sync_turn writes then land in the already-closed old session, and on_session_end tries to commit it a second time — the new session never accumulates messages and never triggers memory extraction. The fix mirrors the pattern Hindsight uses (#17508): 1. Wait for any in-flight sync thread to drain under the OLD _session_id before we mutate it, otherwise the commit below races the last message write. 2. Commit the old session if it accumulated turns — same extraction semantics as on_session_end. Skip if empty (nothing to extract). 3. Drain in-flight prefetch from the old session and clear its cached result so the new session doesn't see stale recall. 4. Rotate _session_id to the new value and reset _turn_count. Commit failures are swallowed (logged at WARN) so a flaky server can't strand the provider on the old session forever — same posture as the existing on_session_end commit.
This commit is contained in:
@@ -420,3 +420,105 @@ def test_viking_client_health_sends_auth_headers(monkeypatch):
|
||||
assert client.health() is True
|
||||
assert captured["url"] == "https://example.com/health"
|
||||
assert captured["headers"]["Authorization"] == "Bearer test-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# on_session_switch — flush + commit + rotate behavior (hermes-agent#28296)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_provider_with_session(session_id: str, turn_count: int):
|
||||
provider = OpenVikingMemoryProvider()
|
||||
provider._client = MagicMock()
|
||||
provider._session_id = session_id
|
||||
provider._turn_count = turn_count
|
||||
return provider
|
||||
|
||||
|
||||
def test_on_session_switch_commits_old_session_and_rotates_id():
|
||||
provider = _make_provider_with_session("old-sid", turn_count=3)
|
||||
|
||||
provider.on_session_switch("new-sid", parent_session_id="old-sid")
|
||||
|
||||
provider._client.post.assert_called_once_with("/api/v1/sessions/old-sid/commit")
|
||||
assert provider._session_id == "new-sid"
|
||||
assert provider._turn_count == 0
|
||||
|
||||
|
||||
def test_on_session_switch_skips_commit_for_empty_old_session():
|
||||
"""No turns accumulated → nothing to extract → no commit call."""
|
||||
provider = _make_provider_with_session("old-sid", turn_count=0)
|
||||
|
||||
provider.on_session_switch("new-sid")
|
||||
|
||||
provider._client.post.assert_not_called()
|
||||
assert provider._session_id == "new-sid"
|
||||
assert provider._turn_count == 0
|
||||
|
||||
|
||||
def test_on_session_switch_clears_stale_prefetch_result():
|
||||
provider = _make_provider_with_session("old-sid", turn_count=1)
|
||||
provider._prefetch_result = "stale recall from old session"
|
||||
|
||||
provider.on_session_switch("new-sid")
|
||||
|
||||
assert provider._prefetch_result == ""
|
||||
|
||||
|
||||
def test_on_session_switch_waits_for_inflight_sync_thread():
|
||||
"""In-flight sync_turn write must drain before the commit fires —
|
||||
otherwise the commit can race the last message write."""
|
||||
provider = _make_provider_with_session("old-sid", turn_count=2)
|
||||
|
||||
join_calls = []
|
||||
|
||||
class FakeThread:
|
||||
def is_alive(self):
|
||||
return True
|
||||
|
||||
def join(self, timeout=None):
|
||||
join_calls.append(timeout)
|
||||
|
||||
provider._sync_thread = FakeThread()
|
||||
|
||||
provider.on_session_switch("new-sid")
|
||||
|
||||
assert join_calls, "expected on_session_switch to join the in-flight sync thread"
|
||||
provider._client.post.assert_called_once_with("/api/v1/sessions/old-sid/commit")
|
||||
|
||||
|
||||
def test_on_session_switch_noop_on_empty_new_id():
|
||||
provider = _make_provider_with_session("old-sid", turn_count=5)
|
||||
|
||||
provider.on_session_switch("")
|
||||
provider.on_session_switch(" ")
|
||||
|
||||
provider._client.post.assert_not_called()
|
||||
assert provider._session_id == "old-sid"
|
||||
assert provider._turn_count == 5
|
||||
|
||||
|
||||
def test_on_session_switch_noop_when_client_missing():
|
||||
provider = OpenVikingMemoryProvider()
|
||||
provider._client = None
|
||||
provider._session_id = "old-sid"
|
||||
provider._turn_count = 4
|
||||
|
||||
# Must not raise even though no client is configured.
|
||||
provider.on_session_switch("new-sid")
|
||||
|
||||
# State stays untouched — provider is effectively disabled.
|
||||
assert provider._session_id == "old-sid"
|
||||
assert provider._turn_count == 4
|
||||
|
||||
|
||||
def test_on_session_switch_swallows_commit_failure():
|
||||
"""Commit-on-switch must not propagate exceptions: a failing commit on the
|
||||
old session must still allow the rotate to the new session to complete,
|
||||
otherwise subsequent sync_turn writes would land in the wrong session."""
|
||||
provider = _make_provider_with_session("old-sid", turn_count=2)
|
||||
provider._client.post.side_effect = RuntimeError("commit boom")
|
||||
|
||||
provider.on_session_switch("new-sid")
|
||||
|
||||
assert provider._session_id == "new-sid"
|
||||
assert provider._turn_count == 0
|
||||
|
||||
Reference in New Issue
Block a user