cowork-local/ui/cowork_tab.py

"""Cowork tab — chat with the local Internal Agent."""
from __future__ import annotations

from pathlib import Path

from PySide6.QtWidgets import (
    QFileDialog, QHBoxLayout, QLabel, QPushButton, QVBoxLayout, QWidget,
)

from ..config import PROVIDER_LABELS
from ..core import agent_roles
from ..core.worker import AgentWorker
from ..i18n import on_language_changed, tr
from ..state import AppContext
from .chat_panel import ChatPanel
from .icons import icon

_FOLDER_LBL_MAX_CHARS = 42  # keep the composer's bottom row from crowding out Agent/Send


class CoworkTab(ChatPanel):
    def __init__(self, ctx: AppContext):
        super().__init__(ctx, "cowork", "Cowork", placeholder_key="composer.placeholder_cowork")

        self._title_lbl = QLabel()
        self._title_lbl.setStyleSheet("font-weight:700; font-size:15px;")
        self.model_lbl = QLabel("")
        self.model_lbl.setObjectName("hint")

        self.skills_btn = QPushButton()
        self.skills_btn.setIcon(icon("book"))
        self.skills_btn.clicked.connect(self._open_skills_manager)

        self._new_btn = QPushButton()
        self._new_btn.setIcon(icon("new"))
        self._new_btn.clicked.connect(self.new_session)

        self.toolbar_layout.addWidget(self._title_lbl)
        self.toolbar_layout.addWidget(self.model_lbl)
        self.toolbar_layout.addStretch(1)
        self.toolbar_layout.addWidget(self.skills_btn)
        self.toolbar_layout.addWidget(self._new_btn)

        # Where Cowork saves its deliverables — the project's sandbox workspace
        # (or, for a folder picked via the button, that folder). Only the
        # folder NAME shows beside the button (not the full path as a
        # link/hint text) — the full path is still available as a tooltip.
        self.folder_btn = QPushButton()
        self.folder_btn.setIcon(icon("folder"))
        self.folder_btn.clicked.connect(self._pick_output_folder)
        self.folder_lbl = QLabel("")
        self.folder_lbl.setObjectName("hint")
        folder_box = QWidget()
        _fbl = QHBoxLayout(folder_box)
        _fbl.setContentsMargins(0, 0, 0, 0)
        _fbl.setSpacing(6)
        _fbl.addWidget(self.folder_btn)
        _fbl.addWidget(self.folder_lbl)   # folder name BESIDE the button
        # Which project (workspace) the CURRENT thread belongs to. Cowork now
        # always lives INSIDE the Workspace screen for a single selected project,
        # so the project name beside the composer is redundant chrome — the label
        # is kept as a hidden member (its text still tracks the active project for
        # tooltips/tests) but no longer added to the layout.
        self.project_lbl = QLabel("")
        self.project_lbl.setObjectName("hint")
        self.project_lbl.setVisible(False)
        self.composer.add_bottom_left(folder_box)

        # Per-workspace "Auto-run" toggle (auto-approve commands in THIS
        # workspace) — sits next to the routing toggle the base class added.
        from .routing_toggle import AutoRunToggle
        self.autorun_toggle = AutoRunToggle(self.ctx)
        self.composer.add_bottom_right(self.autorun_toggle)

        self.refresh_header()
        # Start watching the output folder for new files
        wd = self.workspace_dir()
        if wd:
            self._start_watching(wd)
        # Cowork shows ONLY the final deliverables, refreshed from disk once a turn
        # succeeds (see _cleanup_turn) — never intermediate files or a folder name.
        on_language_changed(self._retranslate)

    def refresh_title(self) -> None:
        """Head the screen with the thread you are in, as the drawing does.

        It said "Cowork" on every conversation — the screen's own name, which
        the rail already shows. The thread's title is the thing that changes and
        the thing that tells you where you are; a thread with no title yet (a
        new chat, before its first turn) falls back to the screen name.
        """
        lbl = getattr(self, "_title_lbl", None)
        if lbl is None:
            return          # ChatPanel.__init__ sets self.title before we exist
        lbl.setText(getattr(self, "title", "") or tr("cowork.title"))

    def _retranslate(self) -> None:
        self.refresh_title()
        self.skills_btn.setText(tr("cowork.skills_btn"))
        self.skills_btn.setToolTip(tr("cowork.skills_tooltip"))
        self._new_btn.setText(tr("cowork.new_chat"))
        self.folder_btn.setText(tr("cowork.pick_folder_btn"))
        self.folder_btn.setToolTip(tr("cowork.pick_folder_tooltip"))
        self._apply_output_folder_label()

    # ---- output folder --------------------------------------------------
    def _pick_output_folder(self) -> None:
        start = str(self._session_output_dir())
        chosen = QFileDialog.getExistingDirectory(self, tr("cowork.pick_folder_title"), start)
        if not chosen:
            return
        if self.project_id not in ("", "default"):
            # Inside a project, the picked folder becomes THAT project's
            # workspace (its sandbox + shared-knowledge root) — not the
            # global default-output setting.
            from ..core.projects import save_project

            project = self._project()
            if project is not None:
                project.output_dir = chosen
                save_project(project)
        else:
            self.ctx.config.cowork["output_dir"] = chosen
            self.ctx.save()
        self._apply_output_folder_label()
        self._refresh_outputs_from_disk()

    def _apply_output_folder_label(self) -> None:
        full = str(self._session_output_dir())
        # Only the folder NAME is shown beside the button — the full path
        # (still available on hover) reads as noisy clutter for a value the
        # user just picked and already knows the location of.
        name = Path(full).name or full
        if len(name) > _FOLDER_LBL_MAX_CHARS:
            name = name[:_FOLDER_LBL_MAX_CHARS - 1] + "…"
        self.folder_lbl.setText(name)
        self.folder_lbl.setToolTip(full)
        # Text still tracks the active project (tooltip/tests read it), but the
        # label stays hidden — see its creation note above.
        project = self._project()
        if project is not None:
            self.project_lbl.setText(tr("cowork.project_label", name=project.name))
            self.project_lbl.setToolTip(tr("cowork.project_tooltip", name=project.name))
        else:
            self.project_lbl.setText("")
        # Start watching the output folder for new files
        wd = self.workspace_dir()
        if wd:
            self._start_watching(wd)

    # ---- project (workspace) --------------------------------------------
    def _project(self):
        from ..core.projects import load_project

        return load_project(self.project_id)

    def set_project(self, project_id: str) -> None:
        """Assign the CURRENT thread to a project (called by the Workspace
        screen's 'New chat in project'). Output folder + shared context follow.

        Also makes this the ACTIVE workspace for per-workspace mode resolution
        and refreshes the routing/auto-run toggles to show THIS workspace's
        modes (so switching projects switches the visible modes)."""
        self.project_id = project_id or "default"
        self.ctx.active_project_id = self.project_id
        for t in (getattr(self, "routing_toggle", None), getattr(self, "autorun_toggle", None)):
            if t is not None:
                t.refresh()
        self._apply_output_folder_label()
        self._refresh_outputs_from_disk()

    def project_knowledge_dir(self):
        """A non-default project's shared-knowledge folder — its workspace
        ROOT. This is now the SAME folder every thread of the project already
        saves its deliverables into (no more per-session sub-folder — see
        _session_output_dir), so ChatPanel._augment's equality check skips the
        separate '[Project files]' scan and the '[Workspace files]' scan alone
        already covers everything in one shared place, Claude-Projects style."""
        if self.project_id in ("", "default"):
            return None
        project = self._project()
        return project.workspace_dir() if project else None

    # Generator/helper scripts are never a Cowork deliverable — keep them out of
    # the Output list entirely (only the final file is shown).
    _INTERMEDIATE_EXTS = {".py", ".pyw", ".js", ".mjs", ".cjs", ".ts",
                          ".sh", ".bat", ".ps1", ".rb", ".pl"}

    def assistant_title(self) -> str:
        return tr("cowork.assistant_title")

    def _session_output_dir(self):
        """Where this session's deliverables are saved.

        Default: a sub-folder of the configured Output root named by the
        session id (a timestamp — never the chat content), so same-named
        outputs from different sessions don't overwrite each other.

        Once a folder has been EXPLICITLY specified — either "Chọn thư mục
        khác" (see _pick_output_folder) for the default project, or a
        project's own workspace folder — files are saved DIRECTLY into it
        instead, no auto-created per-session sub-folder: the Output box/link
        always points at exactly that folder, and every thread that shares it
        sees the same files (deliverables AND project knowledge together,
        Claude-Projects style). Running turns still get isolated '.turns/<id>'
        sandboxes (unique across every session in this tab — see
        _turn_output_dir/_turn_seq — moved up on success by
        _promote_turn_outputs, which de-dupes by name via _unique_path), so
        concurrent turns/threads sharing one folder never clash.

        Threads of a NON-default project are sandboxed inside that project's
        own workspace: the agent's tools are confined there (ToolContext
        rejects any path escape), and it's a single shared folder for the
        whole project — not one sub-folder per thread."""
        if self.project_id not in ("", "default"):
            project = self._project()
            if project is not None:
                return project.workspace_dir()
        custom = (self.ctx.config.cowork.get("output_dir") or "").strip()
        if custom:
            return Path(custom).expanduser()
        return self.ctx.config.cowork_output_dir() / self.session_id

    def workspace_dir(self):
        return self._session_output_dir()

    def _turn_output_dir(self, turn_id: str):
        """Each running turn writes into its own '.turns/<id>' sandbox so parallel
        turns never clobber each other's files (run_cowork's cleanup diffs the
        folder before/after, which would misfire on a shared dir). On success the
        finished turn's deliverables are moved up to the session root — see
        _cleanup_turn. The dot-prefixed folder is ignored by the Output list
        (which shows only top-level files)."""
        return self._session_output_dir() / ".turns" / turn_id

    def _is_intermediate_output(self, path: str) -> bool:
        from pathlib import Path
        return Path(path).suffix.lower() in self._INTERMEDIATE_EXTS

    def _existing_output_names(self):
        """Deliverables already sitting in this session's output folder (from
        earlier, already-finished turns) — the current turn writes into its own
        '.turns/<id>' sandbox, so this never includes its own in-flight files."""
        try:
            return sorted(p.name for p in self._session_output_dir().iterdir()
                          if p.is_file() and not p.name.startswith(".")
                          and not self._is_intermediate_output(str(p)))
        except OSError:
            return []

    def _session_notes(self) -> str:
        # Lets the agent (and the user, without re-uploading) reference/revise a
        # file it made earlier in this same conversation — e.g. "sửa lại tiêu đề
        # trong file báo cáo vừa tạo" — since it can read_file/edit_file/write_file
        # it directly by the exact name listed here.
        names = self._existing_output_names()
        if not names:
            return ""
        listing = ", ".join(names)
        return (
            "[Session context] Files already created earlier in this conversation's output "
            "folder — read/revise them directly by their exact name with read_file/edit_file/"
            "write_file; the user does NOT need to re-upload them if they refer to \"the file "
            f"I made\" / \"file vừa tạo\" or similar: {listing}"
        )

    def _promote_turn_outputs(self, turn_dir, record, session_root) -> None:
        """Move a finished turn's files from its '.turns/<id>' sandbox up to its
        conversation's Output root (``session_root`` — the turn's HOME session, so a
        background turn lands in the right chat even after the user switched away),
        then remove the sandbox. run_cowork already flattened and stripped
        '.scratch'/generator scripts, so the sandbox root holds the final file(s).
        Every top-level file is moved (no extension filtering) so a file the user
        genuinely asked for is never dropped before the sandbox is removed — the
        Output list itself decides what to *show* (see _refresh_outputs_from_disk)."""
        import shutil
        from pathlib import Path

        from ..core.chat_agent import _unique_path

        turn_dir = Path(turn_dir)
        session_root = Path(session_root)
        try:
            if not turn_dir.is_dir() or turn_dir.resolve() == session_root.resolve():
                return
        except OSError:
            return
        session_root.mkdir(parents=True, exist_ok=True)
        remap = {}
        try:
            files = [p for p in turn_dir.iterdir() if p.is_file()]
        except OSError:
            files = []
        for p in files:
            try:
                dest = _unique_path(session_root, p.name)
                p.replace(dest)
                remap[str(p)] = str(dest)
            except OSError:
                pass
        shutil.rmtree(turn_dir, ignore_errors=True)
        # Keep this turn's recorded outputs pointing at the moved files so
        # deleting the message later still finds and removes them.
        if remap and record is not None:
            record["outputs"] = [remap.get(p, p) for p in record.get("outputs", [])]

    # ---- Output box: only the final, successful deliverables ----------
    def register_output(self, path: str) -> None:
        # Suppress live/intermediate updates during a run — the Output box is
        # rebuilt from the surviving files once the turn succeeds (see below).
        return

    def on_file_written(self, path: str) -> None:
        return  # no live file updates in Cowork

    def _refresh_outputs_from_disk(self) -> None:
        """Show only the final deliverable files now sitting in the session folder
        (no scripts/intermediate files, no hidden/`.scratch`, no folders)."""
        self.output_section.clear()
        try:
            files = sorted((p for p in self._session_output_dir().iterdir() if p.is_file()),
                           key=lambda p: p.name)
        except OSError:
            return
        for p in files:
            if not p.name.startswith(".") and not self._is_intermediate_output(str(p)):
                self.output_section.add(str(p))
        self.output_changed.emit(str(self._session_output_dir()))

    def new_session(self) -> None:
        # The new thread stays in the CURRENT project (Claude-style).
        super().new_session()
        # Refresh the project/folder labels + watch the new session's folder.
        self._apply_output_folder_label()

    def load_conversation(self, conv) -> None:
        super().load_conversation(conv)    # restores this conversation's project_id
        self._refresh_outputs_from_disk()  # show this session's deliverables from disk
        # Refresh the project/folder labels + watch this conversation's folder.
        self._apply_output_folder_label()

    def refresh_header(self) -> None:
        cfg = self.ctx.config
        label = PROVIDER_LABELS.get(cfg.active_provider, cfg.active_provider)
        self.model_lbl.setText(f"{label} · {cfg.model_label()}")
        self._apply_output_folder_label()  # picks up edits made via Settings too

    def build_job(self, text: str, messages, out_dir):
        """Build the worker job for one Cowork turn (R04-T04).

        Every input the turn needs is captured HERE, on the UI thread, into an
        immutable ``ConversationExecutionRequest``. Previously the job closure
        read widget state (selected model, active workspace, project
        instructions) from inside the worker thread, so a turn could run on a
        mixture of the state at submit time and the state the user changed while
        it was running - and which mixture you got depended on thread timing.

        Each turn still writes into its OWN isolated folder (out_dir) and works
        on its OWN message list, so several turns can run in parallel without
        clobbering each other's files or history. Deliverables are moved up to
        the session Output root when the turn finishes (see _cleanup_turn).
        """
        from ..core.projects import load_project, project_context_text
        from ..domain.agents import ConversationExecutionRequest

        output_dir = out_dir or self._session_output_dir()
        # Shared project instructions (Claude-Projects style), read now so edits
        # made in the Workspace screen mid-turn cannot change this turn's prompt.
        project_context = project_context_text(load_project(self.project_id))
        # The Admin-defined agent preset's instructions, if one is selected.
        agent_prompt = self.admin_agent_prompt()
        if agent_prompt:
            project_context = (f"{project_context}\n\n{agent_prompt}"
                               if project_context else agent_prompt)

        request = ConversationExecutionRequest.create(
            text, messages,
            output_dir=str(output_dir),
            session_id=self.session_id,
            surface=self.kind,
            title=self.title,
            project_id=self.project_id,
            project_context=project_context,
            agent_role=agent_roles.COWORK,
            # Permission Management (Sandbox Security Layer): off by default -
            # matches the pre-existing auto-run behavior. Resolved PER WORKSPACE:
            # this project's Auto-run override wins, else the global setting.
            confirm_commands=bool(self.ctx.project_confirm_commands()),
        )
        # Built on the UI thread with everything else: it already reflects this
        # turn's routing decision and the selected agent/model.
        provider = self.build_provider()

        def job(worker: AgentWorker):
            from ..application.conversations import ConversationApplicationService

            service = ConversationApplicationService(
                # The provider is part of the snapshot, so the factory ignores
                # the request's provider/model rather than re-resolving them
                # from live config inside the worker thread.
                lambda _provider_id, _model: provider,
                # 🔌 MCP Layer: every tool source flows through MCP now - the
                # external servers configured in Settings AND Microsoft 365 (a
                # built-in MCP server auto-registered while signed in, see
                # AppContext._ms365_builtin_connection / mcp_servers/ms365_server.py).
                tool_source=self.ctx.build_mcp_tools,
                gate_factory=lambda req: worker.new_gate("confirm",
                                                         agent_role=agent_roles.COWORK),
                security_config=self.ctx.config,
            )
            result = service.run_turn(
                request,
                # The typed events are rendered back into the legacy dict shape
                # the chat widgets already consume; they migrate to AgentEvent
                # directly in EPIC R08.
                on_event=lambda event: worker.emit_event(event.to_dict()),
                cancel=worker.is_cancelled,
            )
            # The service reports a failure instead of raising, but this worker's
            # contract is exception-based (core/worker.py turns one into the
            # `failed` signal that _on_failed already handles), so re-raise the
            # ORIGINAL exception to keep that path byte-for-byte unchanged.
            result.raise_if_failed()
            return {"messages": result.messages, "turn_dir": str(output_dir)}

        return job

    def _cleanup_turn(self, ctx, ok) -> None:
        """A turn ended (successfully or not): promote whatever files it
        produced up to ITS conversation's Output root, THEN discard the
        (by then empty, or intermediate-only) sandbox.

        This runs the SAME promotion on failure as on success — a turn can
        genuinely create a deliverable (e.g. save_file succeeds) and THEN hit
        an unrelated error later in the same turn (a follow-up tool call, a
        network drop, a rate limit that didn't recover) which raises and
        marks the whole turn as failed. Discarding the sandbox unconditionally
        in that case would silently delete a file the user actually got —
        exactly the "file đã tạo bị xóa" bug. Promoting first is always safe:
        an empty/intermediate-only sandbox just promotes zero files.

        Refresh the visible Output list only when that conversation is the
        one on screen."""
        turn_dir = ctx.get("out_dir")
        home_root = ctx.get("home_out_root")
        live = ctx.get("home_id") == self.session_id and not ctx.get("detached")
        if turn_dir and home_root:
            self._promote_turn_outputs(turn_dir, ctx.get("record"), home_root)
        elif turn_dir:
            import shutil
            shutil.rmtree(turn_dir, ignore_errors=True)
        if live:
            self._refresh_outputs_from_disk()