feat: headless Mongo session restore for remote AI chat

[ocervell]

Repo 1/3 of Workspace AI Assistant (secator → secator-api → secator-ui).

The Mongo-channel chat feature spawns the ai task in mode="chat", interactive="remote"; it talks back and forth through the workspace MongoDB (_type:"ai" findings), and on respawn (after a wait-timeout) the next user message restores the conversation from Mongo. The channel (RemoteBackend), the Ai output type, and the timeout already existed. This PR closes the two remaining core gaps: headless Mongo session restore, and verifying the dispatched worker's query-engine resolves to Mongo.

History-fidelity finding (the Step-0 gate)

Verdict: text-only restore — sufficient for mode="chat" v1 (proceeded).

The litellm ChatHistory carries system / user / assistant / assistant-with-tool_calls / tool-result messages. But the persisted _type:"ai" docs only carry text turns — ai_type="prompt"→user, ai_type="response"→assistant — plus action display records (task/shell/query/…). The assistant tool_calls messages and their tool results are added to the in-memory history but are never yielded as _type:"ai" docs, so they are not in Mongo.

Therefore a literal restore can only rebuild text turns. Crucially, fabricating an assistant tool_calls message without its matching tool result yields a malformed transcript that most providers reject — so we deliberately collapse tool activity into the surrounding text turns rather than reconstruct it. This is:

• Sufficient for mode="chat" (workspace Q&A — the feature's primary mode): the assistant's response text already narrates what it did, so the LLM continues coherently.
• Lossy for mode="attack" (tool-heavy): intermediate tool I/O isn't replayed. Documented follow-up: persist a richer assistant record (the full tool_calls/tool messages) to Mongo if attack-mode replay is needed. This matches the build-spec's flagged risk.

Per the task's own guidance ("a text-only restore that works for mode=\"chat\" is an acceptable v1 if you document the limitation"), I proceeded.

Restore approach

restore_history_from_db(session_id, query_engine, model, encryptor, system_prompt) (secator/ai/session.py):

• query_engine.search({"_type":"ai","session_id":session_id}), sorted by _timestamp.
• prompt→add_user, response→add_assistant; all other ai_types (action displays, follow_up/permission, shell_output, summaries) are channel/UX artifacts and skipped.
• System prompt re-set; turns re-encrypted via maybe_encrypt when an encryptor is active (persisted docs hold plaintext, since response content is decrypted before it's yielded).
• Backend/search errors are swallowed (a respawn must not crash the worker) and return the system-only history.

Wired in ai.py:yielder via _maybe_resume_remote(): when interactive=="remote" and the session has prior _type:"ai" docs, restore from Mongo, append the new user message, and run the loop; fresh conversations (no docs) start as before. The local CLI replay_session/show_session_picker path is untouched. session_id now prefers run_opts.context.session_id (the UI's stable id, reused on respawn). Local history.json writes are skipped on the remote path via a _save_history() helper (Mongo is the source of truth).

Query-engine finding (Task 2)

ctx.get_query_engine() builds QueryEngine(workspace_id, context=self.context); QueryEngine._select_backend resolves the backend from context['drivers'] (mongodb > api > local). The API appends the mongodb driver on every dispatch, so for an async-dispatched ai task with context.workspace_id this correctly resolves to the Mongo backend — no wiring fix needed. The gap was the absence of a guard, so _maybe_resume_remote now emits a Warning when interactive="remote" but the resolved backend is not mongodb/api (the answer channel can't work otherwise). The existing _run_loop RemoteBackend wiring (self.backend.query_engine = ctx.get_query_engine()) uses the same context and is consistent.

Tests (tests/unit/test_ai_session.py, 9 passing)

• restore_history_from_db: rebuilds equivalent History (timestamp order, roles, system prompt, non-turn docs skipped); no-prior-docs → system-only / empty; empty-content docs skipped; backend search failure → system-only (no crash); encryptor re-encrypts restored turns.
• Remote-resume branch: picks Mongo restore when prior docs exist, falls through to fresh start otherwise, and warns on a non-Mongo backend.

Verification

• flake8 (max-line-length 120) clean on all changed files.
• New tests: 9 passed. Related AI suites (test_ai_session, test_ai_interactivity, test_ai_history, test_ai_task_opts): 73 passed. Imports clean.
• Pre-existing unrelated failures in test_ai_loop.py (guardrails "command not approved") reproduce on clean origin/main — not caused by this change.

🤖 Generated with Claude Code

https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm

Summary by CodeRabbit

• New Features
  • Remote interactive sessions can now resume from previously saved chat history.
  • Restored conversations keep the original message order and can include an initial system prompt.
• Bug Fixes
  • Improved session continuity by keeping messages tied to a stable session ID.
  • Prevented malformed or non-conversational items from appearing in restored transcripts.
  • Added safer fallback behavior when session history can’t be read.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5a9b8d91-31bb-458d-ba54-55e30cd32701

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

Walkthrough

Adds remote AI session resumption by introducing restore_history_from_db in session.py to reconstruct ChatHistory from MongoDB workspace documents. The ai task gains _maybe_resume_remote, _get_query_engine, and _save_history helpers; session identity derivation is stabilized; all _run_loop exit paths are unified to _save_history; and unit tests cover the new paths.

Changes

Remote AI Session Resume

Layer / File(s)	Summary
ChatHistory reconstruction from MongoDB secator/ai/session.py	Adds restore_history_from_db: queries _type:"ai" docs by session_id, sorts by _timestamp, maps prompt/response to user/assistant turns in a ChatHistory, optionally sets a system prompt, re-encrypts via encryptor, and returns gracefully on query failure.
Remote resume helpers, yielder stamping, and session identity secator/tasks/ai.py	Imports save_history; stabilizes session_id derivation in _init_options to prefer run_opts.context["session_id"]; stamps Ai items missing session_id inside yielder(); adds an interactive == "remote" early-exit branch; introduces _get_query_engine, _maybe_resume_remote (backend validation, prior-doc search, history rebuild, run loop continuation), and _save_history (skips local history.json writes on remote path).
_run_loop persistence unified to _save_history secator/tasks/ai.py	Replaces all direct save_history(...) call sites in _run_loop (empty-response stop, stop-tool, None redetect, KeyboardInterrupt, auth/connectivity error, generic exception) with the new _save_history() helper.
Unit tests for restore and remote resume tests/unit/test_ai_session.py	Adds TestRestoreHistoryFromDB (ordering, role mapping, system prompt, empty content skip, exception fallback, encryptor re-encryption) and TestRemoteResumeBranch (no-docs → False, Mongo docs → True with correct invocations, non-Mongo backend → Warning emission).

Sequence Diagram(s)

sequenceDiagram
  rect rgba(100, 149, 237, 0.5)
    Note over ai_task,_run_loop: Remote Respawn Path
    ai_task->>_maybe_resume_remote: interactive == "remote", session_id
    _maybe_resume_remote->>_get_query_engine: workspace config
    _get_query_engine-->>_maybe_resume_remote: QueryEngine
    _maybe_resume_remote->>QueryEngine: search(_type="ai", session_id=session_id)
    QueryEngine-->>_maybe_resume_remote: prior_docs[]
  end
  alt No prior docs
    _maybe_resume_remote-->>ai_task: restored=False (fresh start)
  else Non-Mongo backend
    _maybe_resume_remote-->>ai_task: yield Warning("remote")
  else Prior docs found (MongoDB)
    _maybe_resume_remote->>restore_history_from_db: session_id, engine, model, encryptor, system_prompt
    restore_history_from_db->>ChatHistory: rebuild turns from sorted docs
    restore_history_from_db-->>_maybe_resume_remote: ChatHistory
    _maybe_resume_remote->>_run_loop: continue with restored history + new prompt
    _run_loop->>_save_history: persist on every exit path
    _maybe_resume_remote-->>ai_task: restored=True
  end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐇 Hoppity-hop through the Mongo store,
I dig up old sessions from sessions before!
Each prompt and response, sorted just right,
The history restored in the blink of a night.
No turn left behind, no context astray —
The rabbit remembers what you meant to say! 🥕

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely summarizes the main change: restoring remote AI chat sessions from Mongo.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/ai-chat

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted.

Error details

{}

[coderabbitai]

Actionable comments posted: 4

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@secator/tasks/ai.py`:
- Around line 561-565: The remote correlation fallback in the task session setup
is too broad because `session_name` can collide across unrelated conversations.
Update the session initialization logic in the AI task flow so
`_maybe_resume_remote()` and the `self.session_id` assignment rely only on
`run_opts.context.session_id` for remote/web resumes, and otherwise fall back to
a non-remote unique identifier such as `self.id` instead of `session_name`. Keep
the stable remote session behavior for respawns, but ensure `session_name` is
not used as the Mongo query key in the `AI` task/resume path.
- Around line 276-279: The prompt handling in ai.py is treating user-supplied
prompt text as a worker-local file path, which can leak local files when the
remote resume flow passes a path-like string. Update the prompt resolution logic
in the code that assigns self.prompt from run_opts so it only uses the provided
prompt as literal text, and remove the Path(...).is_file() / read_text()
expansion for remote/UI-originated prompts. Keep any file-loading behavior
confined to explicitly trusted/local paths elsewhere, and make sure the prompt
passed into the transcript/LLM is exactly the user input.
- Around line 286-293: Skip restoring text-only history for non-chat modes,
since restore_history_from_db() drops tool calls/results and breaks resumed
remote sessions. In ai.py, update the initialization path around _detect_mode(),
get_system_prompt(), and restore_history_from_db() so history is rebuilt only
when self.mode is chat (or equivalent text-only mode), and leave remote modes
like attack to start with a fresh tool context. Ensure _run_loop() receives a
complete context for any mode that relies on tools.

In `@tests/unit/test_ai_session.py`:
- Line 126: The tests in test_ai_session are unpacking task and engine from
_make_task without using both values, triggering RUF059 and weakening the
assertion intent. Update the affected test cases to only bind the values that
are actually needed, or use the engine variable to assert the
restore/query-engine call where applicable, especially in the restore-related
test methods. Keep the assertions aligned with the behavior under test by
checking the query engine argument in the relevant restore calls.

🪄 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: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7a0b44ab-24bf-4d04-9039-73d5ac61f9da

📥 Commits

Reviewing files that changed from the base of the PR and between 2164993 and 3a8f790.

📒 Files selected for processing (3)

• secator/ai/session.py
• secator/tasks/ai.py
• tests/unit/test_ai_session.py