cowork-local/presentation/chat/chat_session_store.py

"""Lưu, nạp lại phiên chat và đếm token — R08-T06.

``_reattach_running_turn`` là phần tinh tế nhất: người dùng chuyển sang
phiên khác rồi quay lại trong khi lượt cũ vẫn đang chạy, thì phải nối
lại đúng luồng đó chứ không được khởi động lại.

``_compress_messages`` nén ngữ cảnh khi hội thoại dài quá cửa sổ model.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional
from PySide6.QtCore import Qt
from PySide6.QtWidgets import QMessageBox
from ...core.worker import AgentWorker
from ...i18n import tr


class ChatSessionMixin:
    """Trộn vào ChatPanel."""

    def _save_snapshot(self, session_id: str, messages: List[Dict[str, Any]],
                       title: str, inputs: Optional[List[str]] = None,
                       history_dir: Optional[Path] = None) -> None:
        """Persist a conversation by id (used both to register it in History the
        moment it starts and to save a finished background turn). No-op until it has
        a user message. Never raises into the UI.

        ``history_dir``, when given, is used INSTEAD of
        ``self.ctx.config.history_dir()`` — see ``_persist_session`` (R06-T04):
        a background turn must save into the project it started in, not
        whichever project happens to be selected in the Workspace screen by
        the time the turn finishes.
        """
        if not self.ctx.config.history.get("autosave", True):
            return
        if not any(m.get("role") == "user" for m in messages):
            return
        try:
            from ...core.history import save_conversation
            save_conversation(
                history_dir if history_dir is not None else self.ctx.config.history_dir(),
                self.kind, session_id,
                messages, title, inputs=list(inputs or []), outputs=[],
                # Only the CURRENT view knows its project for sure; a background
                # turn's save must not overwrite another conversation's project
                # with whatever the user is viewing now (save_conversation keeps
                # the stored value when '' is passed).
                project_id=self.project_id if session_id == self.session_id else "",
            )
        except Exception:
            pass  # persistence must never disrupt the UI

    def _persist_session(self, ctx: Dict[str, Any]) -> None:
        """Save a BACKGROUND turn's conversation (it isn't the current view, so the
        view-based _autosave can't). Outputs are rebuilt from disk on reopen."""
        self._save_snapshot(ctx["home_id"], ctx["home_messages"],
                            ctx.get("home_title", ""),
                            inputs=ctx.get("record", {}).get("inputs", []),
                            history_dir=ctx.get("home_history_dir"))
        self.history_changed.emit()

    def running_session_ids(self):
        """Set of conversation ids that currently have a turn running (for the
        History status markers)."""
        return set(self._sessions_live)

    def _usage_label(self) -> str:
        return self.title or self.session_id

    def _session_events(self):
        from ...core import usage_tracker as ut
        label = self._usage_label()
        return [e for e in ut.load_events()
                if e.get("source") == self.kind and e.get("label") == label]

    def refresh_usage(self) -> None:
        """Show what this conversation has already cost.

        The label was written only at the end of a turn, so opening a thread
        from History left the strip blank however much it had spent.
        """
        from ...core import model_pricing as mp
        from ...core import usage_tracker as ut

        cur = self._usage_snapshot()
        if not (cur["in"] or cur["out"] or cur["cache"]):
            self._usage_total_lbl.setText("")
            return
        # same source _show_usage reads, so the two never disagree
        pricing = {**ut.DEFAULT_PRICING, **(self.ctx.config.data.get("usage") or {})}
        self._usage_total_lbl.setText(
            f"↓{mp.format_tokens(cur['in'])} ↑{mp.format_tokens(cur['out'])} "
            f"▤{mp.format_tokens(cur['in'] + cur['out'] + cur['cache'])} "
            f"{ut.format_cost(self._session_cost_usd(), pricing)}")

    def _usage_snapshot(self) -> Dict[str, int]:
        """Cumulative in/out/cache tokens for THIS conversation so far."""
        snap = {"in": 0, "out": 0, "cache": 0}
        for e in self._session_events():
            snap["in"] += int(e.get("in", 0) or 0)
            snap["out"] += int(e.get("out", 0) or 0)
            snap["cache"] += int(e.get("cache", 0) or 0)
        return snap

    def _session_cost_usd(self) -> float:
        from ...core import model_pricing as mp
        return sum(mp.turn_cost_usd(e.get("model", ""), e.get("in", 0), e.get("out", 0),
                                    self.ctx.config) for e in self._session_events())

    def _show_usage(self, ctx: Dict[str, Any]) -> None:
        """Per-turn footer under the assistant message + the running conversation
        total (bottom-left). Cost uses the Monitoring model-price table and the
        display currency, and auto-updates when the model is switched."""
        from ...core import model_pricing as mp, usage_tracker as ut
        cur = self._usage_snapshot()
        base = ctx.get("usage_base") or {"in": 0, "out": 0, "cache": 0}
        d_in = max(0, cur["in"] - base.get("in", 0))
        d_out = max(0, cur["out"] - base.get("out", 0))
        d_cache = max(0, cur["cache"] - base.get("cache", 0))
        pricing = {**ut.DEFAULT_PRICING, **(self.ctx.config.data.get("usage") or {})}
        # Condensed format (tight icon+value, single-space separators) — the
        # old 4-space-wide separators made this label wide enough that it got
        # crowded out of the composer's bottom row by the Local-folder button
        # sharing the same row.
        bub = ctx.get("last_assistant")
        if bub is not None and (d_in or d_out):
            turn_usd = mp.turn_cost_usd(self._model, d_in, d_out, self.ctx.config)
            bub.add_usage(f"↓{mp.format_tokens(d_in)} ↑{mp.format_tokens(d_out)} "
                          f"▤{mp.format_tokens(d_in + d_out + d_cache)} "
                          f"{ut.format_cost(turn_usd, pricing)}")
        self._usage_total_lbl.setText(
            f"↓{mp.format_tokens(cur['in'])} ↑{mp.format_tokens(cur['out'])} "
            f"▤{mp.format_tokens(cur['in'] + cur['out'] + cur['cache'])} "
            f"{ut.format_cost(self._session_cost_usd(), pricing)}")

    def _autosave(self) -> None:
        if not self.ctx.config.history.get("autosave", True):
            return
        if not any(m.get("role") == "user" for m in self.messages):
            return
        try:
            from ...core.history import save_conversation
            path = save_conversation(
                self.ctx.config.history_dir(), self.kind, self.session_id,
                self.messages, self.title,
                inputs=self.input_section.paths(),
                outputs=self.output_section.paths(),
                project_id=self.project_id,
            )
            # Remember this as the session to restore next launch (crash-safe).
            last = self.ctx.config.data.setdefault("last_session", {})
            if last.get(self.kind) != str(path):
                last[self.kind] = str(path)
                self.ctx.save()
        except Exception:
            pass  # autosave must never disrupt the UI

    def _maybe_notify_teams(self, result: Dict[str, Any]) -> None:
        teams = self.ctx.config.teams
        notifier = self.ctx.teams_notifier()
        if not (teams.get("notify_on_complete") and notifier.configured):
            return
        summary = self._last_assistant_text() or "Task completed."
        facts = {"Session": self.session_name, "Model": self.ctx.config.model_label()}
        wd = self.workspace_dir()
        if wd:
            facts["Folder"] = str(wd)
        if result.get("error"):
            facts["Status"] = "Error"

        def job(worker: AgentWorker):
            ok, detail = notifier.send(f"Cowork {self.session_name} — task done", summary[:1200], facts)
            return {"ok": ok, "detail": detail}

        w = AgentWorker(job)
        w.finished_ok.connect(lambda r: self.status_message.emit(r.get("detail", "")))
        self._teams_worker = w
        w.start()

    def new_session(self) -> None:
        from ...core.history import new_session_id

        # Allowed while work is running: current turns keep going in the background.
        self._detach_live_turns()
        self.messages = []
        self.session_id = new_session_id()
        self.title = ""
        self.turns = []
        self.chat_view.clear()
        self.composer.clear_queue()
        self.composer.reset_input()   # clear leftover text / "Attached: …" hint
        self.plan_section.clear()
        self.input_section.clear()
        self.output_section.clear()
        self.graph_event.emit(self.session_name, {"type": "reset"})
        self._sync_indicators()
        self.history_changed.emit()   # current view changed → refresh History highlight

    def _notify_title(self) -> None:
        """Let a screen that heads itself with the thread title follow along.

        The thread also decides what the usage strip should read, so refresh
        that here rather than at each of the three places the title changes.
        """
        hook = getattr(self, "refresh_title", None)
        if callable(hook):
            hook()
        if getattr(self, "_usage_total_lbl", None) is not None:
            self.refresh_usage()

    def load_conversation(self, conv: Dict[str, Any]) -> None:
        """Switch the view to a stored conversation. Allowed while work is running —
        the current turns keep going in the background."""
        sid = conv.get("session_id") or self.session_id
        # Clicking the conversation you're already viewing while it has a running
        # turn must NOT tear down its live rendering — just no-op.
        if sid == self.session_id and self._view_busy():
            return
        self._detach_live_turns()
        self.session_id = sid
        self.title = conv.get("title", "")
        self._notify_title()
        self.project_id = conv.get("project_id", "") or "default"
        # If this conversation still has a turn running in the background, attach to
        # its LIVE message list (not a stale disk copy) so the two never race on save.
        if sid in self._sessions_live:
            self.messages = self._sessions_live[sid]
        else:
            self.messages = list(conv.get("messages", []))
        self.turns = []
        self.chat_view.clear()
        self.composer.clear_queue()
        self.composer.reset_input()   # clear leftover text / "Attached: …" hint
        self.plan_section.clear()
        self.input_section.clear()
        self.output_section.clear()
        self.graph_event.emit(self.session_name, {"type": "reset"})
        for m in self.messages:
            role = m.get("role")
            if role == "user":
                self.chat_view.add_user(m.get("content", ""))
                self.graph_event.emit(self.session_name, {"type": "user", "content": m.get("content", "")})
            elif role == "assistant":
                if m.get("content"):
                    self.chat_view.add_assistant(self.assistant_title()).set_markdown(m["content"])
                    self.graph_event.emit(self.session_name, {"type": "assistant_done", "content": m["content"]})
                for tc in m.get("tool_calls", []) or []:
                    self.graph_event.emit(self.session_name, {
                        "type": "tool_proposed", "name": tc.get("name", ""),
                        "args": tc.get("arguments", {}),
                        "preview": {"text": str(tc.get("arguments", {}))},
                    })
            elif role == "tool":
                self.chat_view.add_tool(m.get("name", "tool"), m.get("content", ""), True)
                self.graph_event.emit(self.session_name, {
                    "type": "tool_result", "name": m.get("name", ""),
                    "ok": True, "output": m.get("content", ""),
                })
        # Restore the Input/Output file lists too.
        for p in conv.get("inputs", []):
            self.input_section.add(p)
        for p in conv.get("outputs", []):
            self.output_section.add(p)
        # If this conversation has a turn running in the background, re-render the
        # in-progress task and re-attach it so it keeps streaming live here.
        running = self._running_ctx_for(sid)
        if running is not None:
            self._reattach_running_turn(running)
        elif self.messages:
            # A past (already finished) session — surface a link to its output
            # folder even though the live "done" marker isn't replayed.
            folder = self.workspace_dir()
            if folder:
                marker = self.chat_view.add_status(tr("chat.session_folder_marker"))
                marker.add_folder_link(str(folder), tr("chat.open_folder_short"))
        # Jump to the newest message after the transcript is rebuilt.
        self.chat_view.scroll_to_bottom()
        self._sync_indicators()
        self.history_changed.emit()   # current view changed → refresh History highlight

    def _delete_turn(self, turn: Dict[str, Any]) -> None:
        files = [p for p in (turn.get("inputs", []) + turn.get("outputs", [])) if p]
        if files:
            preview = "\n".join("• " + str(p) for p in files[:12])
            prompt = tr("chatpanel.delete_confirm_files", n=len(files), preview=preview)
        else:
            prompt = tr("chatpanel.delete_confirm_plain")
        if QMessageBox.question(self, tr("chatpanel.delete_confirm_title"), prompt) != QMessageBox.Yes:
            return
        for bubble in turn.get("bubbles", []):
            bubble.setParent(None)
            bubble.deleteLater()
        ids = {id(m) for m in turn.get("messages", [])}
        if ids:
            self.messages = [m for m in self.messages if id(m) not in ids]
        for p in files:
            try:
                fp = Path(p)
                if fp.is_file():
                    fp.unlink()
            except OSError:
                pass
        if turn in self.turns:
            self.turns.remove(turn)
        self._rebuild_io()
        self._autosave()
        self.status_message.emit(tr("chatpanel.delete_done"))

    def _compress_messages(self) -> None:
        """Manual compress: keep the system prompt + the last 2 turns verbatim and
        DIGEST all older messages into one compact summary, shrinking it until the
        whole conversation is under 25% of its original token size."""
        if self._view_busy():
            self.status_message.emit(tr("chatpanel.compress_busy"))
            return
        from ...core.usage_tracker import estimate_tokens

        msgs = list(self.messages)

        def _tok(ms):
            return sum(estimate_tokens(str(m.get("content", ""))) for m in ms)

        orig = _tok(msgs)
        systems = [m for m in msgs if m.get("role") == "system"]
        rest = [m for m in msgs if m.get("role") != "system"]
        starts = [i for i, m in enumerate(rest) if m.get("role") == "user"]
        if len(starts) <= 2 or orig <= 0:
            self.status_message.emit(tr("chatpanel.compress_short"))
            return
        cut = starts[-2]                       # keep the last 2 turns verbatim
        old, recent = rest[:cut], rest[cut:]
        old_tok = _tok(old) or 1               # target: digest < 25% of the OLD part

        def _digest(per_msg: int):
            parts = []
            for m in old:
                c = str(m.get("content", "")).strip().replace("\n", " ")
                if c:
                    parts.append(f"- {m.get('role', '')}: {c[:per_msg]}")
            body = "\n".join(parts)
            return {"role": "user",
                    "content": f"[{tr('chatpanel.compress_digest_header', n=len(old))}]\n{body}"}

        per_msg = 240
        digest = _digest(per_msg)
        # shrink the digest until the OLD conversation is under 25% of its size
        while _tok([digest]) > 0.25 * old_tok and per_msg > 20:
            per_msg = max(20, per_msg // 2)
            digest = _digest(per_msg)
        self.messages = systems + [digest] + recent
        pct = int(_tok([digest]) * 100 / old_tok)
        self.status_message.emit(tr("chatpanel.compress_reduced", pct=pct, n=len(old)))

    def _detach_live_turns(self) -> None:
        """Before switching away from the current conversation, turn its running
        turns into background jobs: they stop rendering into the (about-to-be-
        cleared) transcript but keep running and save to their own conversation."""
        for c in self._active.values():
            if c.get("home_id") == self.session_id:
                c["detached"] = True
                c["assistant"] = None   # its bubbles are about to be cleared

    def _running_ctx_for(self, session_id: str) -> Optional[Dict[str, Any]]:
        """The in-progress turn's context for a conversation (one at a time), or None."""
        for c in self._active.values():
            if c.get("home_id") == session_id:
                return c
        return None

    def _reattach_running_turn(self, ctx: Dict[str, Any]) -> None:
        """Re-render an in-progress turn into the current transcript and re-attach it
        so it keeps streaming live — used when reopening a running conversation, so
        the user sees the CURRENT task (message + steps so far + live plan), not just
        the last saved state."""
        record = ctx["record"]
        record["bubbles"] = []   # the old bubbles were cleared on the view switch
        # 1) the user's message that is being processed
        ub = self.chat_view.add_user(ctx.get("display_text") or "(attachment)")
        record["bubbles"].append(ub)
        # 2) steps already completed this turn (assistant text / tool results); found
        #    by identity after the user message (a system prompt may sit before it).
        #    Snapshot the list — the worker thread may still be appending to it.
        msgs = list(ctx.get("messages", []))
        ui = next((i for i, m in enumerate(msgs) if m is ctx.get("user_msg")), -1)
        for m in (msgs[ui + 1:] if ui >= 0 else []):
            role = m.get("role")
            if role == "assistant" and (m.get("content") or "").strip():
                b = self.chat_view.add_assistant(self.assistant_title())
                b.set_markdown(m["content"])
                record["bubbles"].append(b)
            elif role == "tool":
                b = self.chat_view.add_tool(m.get("name", "tool"), m.get("content", ""), True)
                record["bubbles"].append(b)
        # 3) the live plan checklist (if any) — inline, expandable
        steps = ctx.get("plan_steps") or []
        if steps:
            self.on_plan(steps)
            from ...ui.chat_panel import _format_plan_steps
            pb = self.chat_view.add_plan(_format_plan_steps(steps))
            record["bubbles"].append(pb)
            ctx["plan_bubble"] = pb
        # 4) the partial answer of the step currently streaming — re-attach so new
        #    deltas keep appending to this bubble.
        ctx["assistant"] = None
        ctx["reasoning"] = None
        if (ctx.get("partial") or "").strip():
            ab = self.chat_view.add_assistant(self.assistant_title())
            ab.set_markdown(ctx["partial"])
            record["bubbles"].append(ab)
            ctx["assistant"] = ab
        # 5) live again → future events render here
        ctx["detached"] = False
        self.chat_view.scroll_to_bottom()