feat(claude-code): mid-session dataset switch (seal bridge + re-register) (SDK-306)

[rshkarin]

Review + rework of community PR #186 (by @wiz-abhi) for cognee hackathon issue topoteretes/cognee#3689 — mid-session dataset switch (seal bridge + re-register). Tracked as SDK-306. Base branch main (this repo has no dev).

Switching datasets mid-session repoints where new memory is written without restarting Claude and without losing the current conversation: seal the old (dataset, session_id) bridge, re-register the agent on the new dataset, and keep the same session_id so prior-turn recall still works.

The original commit is preserved as the base (attribution intact); maintainer improvements are layered on top as separate single-concern commits.

Critical finding in #186 — the switch didn't actually switch on main

The original Phase 2 persisted the new dataset to the global config.json (plus os.environ in the switch's own short-lived process, plus .cognee/session-config.json if present). None of those reach a subsequent hook on current main:

• load_config() deliberately drops the file's dataset key (_file_excluded = {"dataset"}), so get_dataset(load_config()) never sees the written value.
• load_resolved() carries no dataset, so every write/recall hook falls back to get_dataset(load_config()).
• an os.environ write cannot cross into a sibling hook process, and no .cognee picker reader exists on main (that's the still-draft feat(claude-code): add project-level dataset picker (#3686) #185 / SDK-304).

So after the switch, store-to-session / session-context-lookup / the final sync kept writing to the old/default dataset. The seal and re-register ran and logged, but the feature was a no-op. (Verified with an adversarial multi-agent read of the code, including an empirical load_config() run.)

What changed (maintainer commits)

1. fix: redirect subsequent hooks via the launch map. Persist the active dataset into the existing host-keyed launch map record — the store every hook already reads through resolve_cognee_session_id — and have get_dataset() consult it. Precedence mid-session switch > COGNEE_PLUGIN_DATASET env > config default: the switch is the latest explicit user action and wins even over a launch-time env pin (on main that env var is the only working way to select a non-default dataset, so otherwise the switch would silently no-op). The override is keyed by host_key, so it never leaks into other launches (unlike the machine-global config file). A launch that never switched resolves exactly as before.
2. refactor: drop inert state and scope creep. Removed the switch_state.json ledger and its write-only active/sealed/sealed_at fields, the per-(session,dataset) high-water baseline (a no-op in HTTP mode — per-dataset bridge buckets already partition writes — and in local mode only reachable on the legacy older-cognee cognee.improve fallback), the ignored config.json dataset write, and the .cognee/session-config.json picker write (no reader on main; coupled to unmerged feat(claude-code): add project-level dataset picker (#3686) #185). Seal now just flushes the old dataset's bucket (digest-deduped) and logs old bridge sealed.
3. test: fixture-free rewrite. The original suite used pytest monkeypatch/tmp_path fixtures and couldn't run under python3 tests/test_*.py (the plugin tree's only runner). Rewrote it to the sibling convention (temp-home save/restore, stubbed network, globals()-iterating __main__ runner) and retargeted the assertions at the reworked behavior.
4. docs: align the skill with the launch-map redirect (it still described the removed config/picker writes).

Kept from #186 (correct and required by the issue): the seal in both HTTP and local modes (+ old bridge sealed), the in-place re-register with the same conn_uuid + session_id (+ agent re-registered), and context preservation (the switch never re-mints session_id/conn_uuid, and recall is already session_id-scoped). Net production surface ~517 → ~333 lines.

Verification

• python3 integrations/claude-code/tests/test_dataset_switch.py — 9/9 pass (also under pytest). Full claude-code suite green except the two pre-existing, unrelated urllib-context= failures on main (test_bridge_poll, test_improve_sync).
• ruff check integrations/ + ruff format --check clean (line-length 100) — the sole automated gate for these plugin-tree dirs.
• Cross-process behavior (a fresh process, like every hook): after recording a switch to proj-y, a separate python3 process resolves get_dataset → proj-y; a different launch still resolves the default; a switch wins over a COGNEE_PLUGIN_DATASET pin.

Out of scope / follow-ups

• No-duplicate-writes across a switch is provided in HTTP mode (primary) by the existing per-dataset shadow buckets + dataset-tagged per-turn writes. The local-SDK modern cognee.improve path reads the session cache by session_id, so fully preventing pre-switch turns from being re-emitted into the new dataset depends on cognee core scoping session data by dataset (SDK-283 / cognee PR #4176) — the plugin can't do that client-side without breaking context-preserving recall.
• Repointing the detached session-end final sync (the exit watcher bakes COGNEE_SYNC_DATASET at spawn): interactive hooks + per-turn writes follow the switch immediately; restarting the watcher on the new dataset is a follow-up aligned with the SDK-212 pattern.
• Mirroring to the codex twin.

Closes: rework of #186 · gh cognee#3689 · pairs with the picker #185 / SDK-304 · depends on core SDK-283 for full local-mode dedup.

🤖 Generated with Claude Code