HERMES-GIELIE/tests/test_issue_new_chat_draft_restore.py

"""Regression coverage for New Chat draft-session restoration.

The bug: New Chat -> type draft -> open history -> New Chat created a fresh
empty session instead of returning to the empty session that owns the draft.
"""
from pathlib import Path


ROOT = Path(__file__).parents[1]
SESSIONS_JS = ROOT.joinpath("static", "sessions.js").read_text(encoding="utf-8")
BOOT_JS = ROOT.joinpath("static", "boot.js").read_text(encoding="utf-8")


def _btn_new_chat_handler() -> str:
    start = BOOT_JS.find("$('btnNewChat').onclick=async()=>{")
    end = BOOT_JS.find("$('btnDownload').onclick", start)
    assert start != -1 and end != -1, "btnNewChat handler block not found"
    return BOOT_JS[start:end]


def _load_session_clear_block() -> str:
    start = SESSIONS_JS.find("async function loadSession(")
    clear_start = SESSIONS_JS.find("if (currentSid !== sid || forceReload) {", start)
    clear_end = SESSIONS_JS.find("// Phase 1: Load metadata only", clear_start)
    assert start != -1, "loadSession not found"
    assert clear_start != -1 and clear_end != -1, "loadSession clear block not found"
    return SESSIONS_JS[start:clear_end]


def test_new_session_remembers_regular_empty_session_id():
    start = SESSIONS_JS.find("async function newSession(")
    end = SESSIONS_JS.find("async function loadSession(", start)
    assert start != -1 and end != -1, "newSession block not found"
    body = SESSIONS_JS[start:end]
    assign_idx = body.find("S.session=data.session")
    remember_idx = body.find("_rememberNewChatDraftSession(S.session)")
    assert assign_idx != -1, "newSession must assign S.session from the POST response"
    assert remember_idx > assign_idx, "newSession must remember the created empty session id"
    assert "if(!(options&&options.worktree)) _rememberNewChatDraftSession(S.session);" in body, (
        "worktree-backed new sessions must not become New Chat draft candidates"
    )


def test_new_chat_button_restores_remembered_draft_before_creating_session():
    body = _btn_new_chat_handler()
    restore_idx = body.find("_restoreRememberedNewChatDraftSession")
    new_idx = body.find("await newSession()")
    assert restore_idx != -1, "New Chat button must try the remembered draft session first"
    assert new_idx != -1, "New Chat button must still fall back to creating a session"
    assert restore_idx < new_idx, "draft-session restore must happen before newSession fallback"
    assert "return;" in body[restore_idx:new_idx], (
        "successful draft-session restore must return instead of also creating a new session"
    )


def test_new_chat_empty_reuse_guard_checks_loaded_visible_messages():
    helper_start = BOOT_JS.find("function _currentSessionIsReusableEmptyChat()")
    helper_end = BOOT_JS.find("$('fileInput').onchange", helper_start)
    assert helper_start != -1 and helper_end != -1, "empty-chat reuse helper not found"
    helper = BOOT_JS[helper_start:helper_end]
    assert "S.session.message_count||0" in helper, (
        "empty-chat reuse should still honor the server-side session message count"
    )
    assert "S.messages.some(m=>m&&m.role&&m.role!=='tool')" in helper, (
        "empty-chat reuse must not focus-only when loaded transcript messages are visible"
    )
    assert "S.session.active_stream_id" in helper and "S.session.pending_user_message" in helper, (
        "empty-chat reuse must continue to treat in-flight zero-count sessions as real sessions"
    )

    button_body = _btn_new_chat_handler()
    assert "if(_currentSessionIsReusableEmptyChat())" in button_body
    shortcut_start = BOOT_JS.find("if((e.metaKey||e.ctrlKey)&&e.key==='k')")
    shortcut_end = BOOT_JS.find("// Cmd/Ctrl+, opens/closes Settings", shortcut_start)
    assert shortcut_start != -1 and shortcut_end != -1, "Cmd/Ctrl+K handler block not found"
    shortcut_body = BOOT_JS[shortcut_start:shortcut_end]
    assert "if(_currentSessionIsReusableEmptyChat())" in shortcut_body


def test_restore_helper_validates_candidate_with_session_metadata():
    assert "const NEW_CHAT_DRAFT_SESSION_KEY = 'hermes-new-chat-draft-session';" in SESSIONS_JS
    assert "async function _restoreRememberedNewChatDraftSession()" in SESSIONS_JS
    assert "messages=0&resolve_model=0" in SESSIONS_JS, (
        "helper should validate the hidden zero-message candidate through /api/session metadata"
    )
    assert "_isRestorableNewChatDraftSession(session, true)" in SESSIONS_JS, (
        "candidate must have a non-empty server-side composer_draft before restore"
    )
    assert "await loadSession(sid, {skipLineageResolve:true});" in SESSIONS_JS, (
        "helper should load the exact hidden empty draft session"
    )


def test_session_switch_awaits_immediate_draft_flush_before_loading_target():
    assert "return api('/api/session/draft'" in SESSIONS_JS, (
        "_saveComposerDraftNow should return its POST promise so switch-away can await it"
    )
    assert "await _saveComposerDraftNow(currentSid" in SESSIONS_JS, (
        "loadSession must flush the current draft before fetching the next session"
    )


def test_immediate_empty_draft_flush_clears_locally_known_server_draft():
    assert "const _composerDraftKnownPayloadSessions = new Set();" in SESSIONS_JS
    save_start = SESSIONS_JS.find("function _saveComposerDraft(")
    save_end = SESSIONS_JS.find("function _composerDraftHasPayload", save_start)
    now_start = SESSIONS_JS.find("function _saveComposerDraftNow(")
    now_end = SESSIONS_JS.find("// Restore composer draft", now_start)
    assert save_start != -1 and save_end != -1
    assert now_start != -1 and now_end != -1
    save_body = SESSIONS_JS[save_start:save_end]
    now_body = SESSIONS_JS[now_start:now_end]
    assert "_composerDraftKnownPayloadSessions.add(sid);" in save_body, (
        "debounced non-empty draft saves must mark the session as having server-side draft payload"
    )
    assert "_rememberComposerDraftPayloadState(sid, normalizedText, normalizedFiles);" in save_body
    assert "!_composerDraftKnownPayloadSessions.has(sid)" in now_body, (
        "empty immediate switch-away saves may only be skipped when no local/server draft payload is known"
    )
    assert "_rememberComposerDraftPayloadState(sid, normalizedText, normalizedFiles);" in now_body


def test_pre_switch_draft_flush_rechecks_stale_loading_guard():
    """The awaited draft-save in loadSession yields the event loop. On a rapid
    session switch (B then quickly C) the stale B continuation must bail out
    before the destructive state-clearing block, or it would wipe the
    freshly-loaded C state. The guard is `if (_loadingSessionId !== sid) return;`
    placed AFTER the awaited save and BEFORE the `S.messages = []` clear
    (Codex pre-release CORE catch, #3471)."""
    body = _load_session_clear_block()
    await_idx = body.find("await _saveComposerDraftNow(currentSid")
    guard_idx = body.find("if (_loadingSessionId !== sid) return;", await_idx)
    clear_idx = body.find("S.messages = [];", await_idx)
    assert await_idx != -1, "pre-switch awaited draft save not found"
    assert guard_idx != -1, "stale-loading guard missing after the awaited draft save"
    assert clear_idx != -1, "destructive S.messages clear not found"
    assert await_idx < guard_idx < clear_idx, (
        "the _loadingSessionId stale-guard must sit between the awaited draft "
        "save and the destructive state clear so a rapid switch can't blank the "
        "newer session"
    )


def test_restorable_candidate_rejects_in_flight_worktree_and_cross_profile_sessions():
    start = SESSIONS_JS.find("function _isRestorableNewChatDraftSession(")
    end = SESSIONS_JS.find("function _rememberNewChatDraftSession", start)
    assert start != -1 and end != -1, "restorable candidate helper not found"
    body = SESSIONS_JS[start:end]
    assert "messageCount !== 0" in body
    assert "session.active_stream_id || session.pending_user_message || session.worktree_path" in body
    assert "_profileMatchesActiveProfile(sessionProfile, activeProfile)" in body
    assert "session.composer_draft || {}" in body
    assert "text || files.length" in body


def test_clear_composer_draft_forgets_same_new_chat_candidate():
    start = SESSIONS_JS.find("function _clearComposerDraft(")
    end = SESSIONS_JS.find("const SESSION_VIEWED_COUNTS_KEY", start)
    assert start != -1 and end != -1, "_clearComposerDraft block not found"
    body = SESSIONS_JS[start:end]
    assert "_clearRememberedNewChatDraftSession(sid);" in body, (
        "sending a draft must stop New Chat from restoring that now-cleared candidate"
    )
    assert "return api('/api/session/draft'" in body, (
        "clear path should return its POST promise for callers/tests that need to await it"
    )
    assert "_rememberComposerDraftPayloadState(sid, '', []);" in body, (
        "clear path must also clear local draft-payload tracking so send-then-switch can skip redundant empty flushes"
    )


def test_boot_restore_preserves_zero_message_session_with_composer_draft():
    """A hard refresh should not discard a zero-message session that owns unsent draft text/files."""
    marker = "const _restoredInFlight = S.session && ("
    start = BOOT_JS.find(marker)
    end = BOOT_JS.find("// Restore the panel from localStorage", start)
    assert start != -1 and end != -1, "boot restored-session cleanup block not found"
    body = BOOT_JS[start:end]
    assert "const _restoredDraft = (S.session && S.session.composer_draft) || {};" in body
    assert "const _restoredDraftText = String(_restoredDraft.text||'').trim();" in body
    assert "const _restoredDraftFiles = Array.isArray(_restoredDraft.files)" in body
    assert "const _restoredHasDraft = !!(_restoredDraftText || _restoredDraftFiles.length);" in body
    assert "&& !_restoredInFlight && !_restoredHasDraft" in body, (
        "zero-message restored sessions should only be dropped when they have no draft"
    )