openworker/tests/test_multiroot.py

"""Slice B — multi-root file toolkit + permission scoping + the context injector.

Orphan Cowork sessions own a primary writable scratch dir and may gain additional folders,
each read-only or read-write. These cover the three layers that share the roots list.
"""

from __future__ import annotations

import asyncio
from pathlib import Path

import pytest

import aisuite as ai
from coworker.engine import TurnEngine
from coworker.events import EventType
from coworker.permissions import Decision, Mode, PermissionEngine
from coworker.providers import AssistantTurn, ToolCall
from coworker.roots import RootDir, normalize_roots, render_context
from coworker.tools import ToolRegistry


def _bare_engine(**kw):
    return TurnEngine(
        provider=object(),
        registry=ToolRegistry(),
        permissions=PermissionEngine(workspace_root="/tmp"),
        model="x",
        **kw,
    )


async def _collect(aiter):
    return [e async for e in aiter]


def _roots(tmp_path):
    scratch = tmp_path / "scratch"
    ro = tmp_path / "ro"
    rw = tmp_path / "rw"
    for d in (scratch, ro, rw):
        d.mkdir()
    (ro / "data.txt").write_text("secret", encoding="utf-8")
    return scratch, ro, rw


# -- file toolkit ---------------------------------------------------------------


def test_toolkit_reads_any_root_writes_only_writable(tmp_path):
    scratch, ro, rw = _roots(tmp_path)
    roots = [
        {"path": str(scratch), "writable": True},
        {"path": str(ro), "writable": False},
        {"path": str(rw), "writable": True},
    ]
    tools = {f.__name__: f for f in ai.toolkits.files(roots=roots)}

    # relative path resolves against the primary (scratch)
    assert tools["write_file"]("note.txt", "hi") == "note.txt"
    assert (scratch / "note.txt").read_text() == "hi"
    # read from a read-only root via absolute path
    assert tools["read_file"](str(ro / "data.txt")) == "secret"
    # write into a writable non-primary root via absolute path
    tools["write_file"](str(rw / "added.txt"), "x")
    assert (rw / "added.txt").read_text() == "x"
    # write into a read-only root is refused
    with pytest.raises(PermissionError):
        tools["write_file"](str(ro / "nope.txt"), "x")
    # escaping every root is refused
    with pytest.raises(PermissionError):
        tools["read_file"]("/etc/hosts")


def test_toolkit_runtime_added_root_is_seen_live(tmp_path):
    scratch, ro, rw = _roots(tmp_path)
    shared = normalize_roots([RootDir(path=scratch, writable=True)])
    tools = {f.__name__: f for f in ai.toolkits.files(roots=shared)}
    # initially the ro dir is out of bounds
    with pytest.raises(PermissionError):
        tools["read_file"](str(ro / "data.txt"))
    # add it at runtime by mutating the shared list in place
    shared.append(RootDir(path=ro, writable=False))
    assert tools["read_file"](str(ro / "data.txt")) == "secret"


# -- permission engine ----------------------------------------------------------


def _meta(reg, name):
    return reg.get(name).metadata


def test_permissions_write_blocked_in_readonly_root(tmp_path):
    scratch, ro, rw = _roots(tmp_path)
    reg = ToolRegistry()
    reg.register_all(
        ai.toolkits.files(roots=[{"path": str(scratch), "writable": True}])
    )
    eng = PermissionEngine(
        workspace_root=scratch,
        roots=[
            {"path": str(scratch), "writable": True},
            {"path": str(ro), "writable": False},
        ],
    )
    # write to the read-only root -> denied outright (not even an approval prompt)
    d = eng.evaluate(
        "write_file",
        {"path": str(ro / "x.txt"), "content": "x"},
        _meta(reg, "write_file"),
    )
    assert not d.allowed and not d.needs_user
    # write to the writable scratch -> needs approval (interactive), not denied
    d = eng.evaluate(
        "write_file", {"path": "x.txt", "content": "x"}, _meta(reg, "write_file")
    )
    assert not d.allowed and d.needs_user


# -- context injector -----------------------------------------------------------


def test_render_context_marks_primary_and_access(tmp_path):
    scratch, ro, _ = _roots(tmp_path)
    text = render_context(
        normalize_roots(
            [
                RootDir(path=scratch, writable=True),
                RootDir(path=ro, writable=False),
            ]
        )
    )
    assert "read-write" in text and "read-only" in text
    assert "primary" in text.lower()
    assert str(scratch.resolve()) in text


def test_outbound_messages_appends_context_to_last_user_message():
    eng = TurnEngine(
        provider=object(),
        registry=ToolRegistry(),
        permissions=PermissionEngine(workspace_root="/tmp"),
        model="x",
        messages=[
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi"},
            {"role": "user", "content": "do it"},
        ],
        context_provider=lambda: "DIRS",
    )
    out = eng._outbound_messages()
    # ephemeral: the persisted history is untouched
    assert eng.messages[-1]["content"] == "do it"
    # the block is glued onto the newest user message, framed as automatic context
    assert len(out) == len(eng.messages)
    assert out[-1]["role"] == "user"
    assert out[-1]["content"].startswith("do it\n\n<system-context>\n(automatic per-turn context")
    assert out[-1]["content"].endswith("\nDIRS\n</system-context>")
    assert out[1]["content"] == "hello"


def test_outbound_messages_noop_without_provider():
    eng = TurnEngine(
        provider=object(),
        registry=ToolRegistry(),
        permissions=PermissionEngine(workspace_root="/tmp"),
        model="x",
        messages=[{"role": "user", "content": "hello"}],
    )
    # No context provider → no `<system-context>` block appended (content unchanged). The list is
    # now always a fresh copy (the `source` strip is unconditional), so compare by value not identity.
    out = eng._outbound_messages()
    assert out == eng.messages
    assert out[-1]["content"] == "hello"


# -- Slice C: add/remove session folders (RO/RW) + persistence ------------------


def _cowork_manager(tmp_path):
    from coworker.providers import ModelCapabilities, ProviderClient
    from coworker.server import SessionManager

    class _Provider(ProviderClient):
        def complete(self, *, model, messages, tools=None, **s):
            return AssistantTurn(text="", finish_reason="stop")

        def capabilities(self, model):
            return ModelCapabilities()

    mgr = SessionManager(data_dir=tmp_path / "data", provider=_Provider())
    mgr._prefs["scratch_base"] = str(tmp_path / "scratchbase")
    return mgr


def test_add_and_remove_roots_live_and_persisted(tmp_path):
    mgr = _cowork_manager(tmp_path)
    ro = tmp_path / "shared_ro"
    rw = tmp_path / "shared_rw"
    ro.mkdir()
    rw.mkdir()
    sid = "sessC"
    engine = mgr.get_engine(sid, agent="cowork")
    assert engine is not None

    # only the primary scratch to start
    roots = mgr.get_roots(sid)
    assert len(roots) == 1 and roots[0]["primary"] and roots[0]["writable"]

    # add a read-only and a read-write folder; the live engine sees them immediately
    mgr.add_root(sid, str(ro), writable=False)
    mgr.add_root(sid, str(rw), writable=True)
    assert {r.path for r in engine.roots} == {
        Path(roots[0]["path"]).resolve(),
        ro.resolve(),
        rw.resolve(),
    }
    by_path = {r["path"]: r for r in mgr.get_roots(sid)}
    assert by_path[str(ro.resolve())]["writable"] is False
    assert by_path[str(rw.resolve())]["writable"] is True

    # the permission engine now allows writes into the rw folder but not the ro folder
    assert mgr.get_engine(sid).permissions._under_writable_root(str(rw / "x.txt"))
    assert not mgr.get_engine(sid).permissions._under_writable_root(str(ro / "x.txt"))

    # cannot remove the primary scratch
    assert not mgr.remove_root(sid, roots[0]["path"])["ok"]
    # remove the read-only folder
    mgr.remove_root(sid, str(ro))
    assert ro.resolve() not in {r.path for r in engine.roots}

    # persist (as a turn would) and reload in a fresh manager → the rw folder survives
    mgr.save(sid, engine)
    mgr2 = _cowork_manager(tmp_path)
    persisted = {r["path"]: r for r in mgr2.get_roots(sid)}
    assert (
        str(rw.resolve()) in persisted
        and persisted[str(rw.resolve())]["writable"] is True
    )
    assert str(ro.resolve()) not in persisted


def test_add_root_before_first_turn_persists(tmp_path):
    """Adding a folder on a brand-new conversation (no record, no engine yet) must survive:
    the manager creates a minimal cowork record so the grant isn't lost (GUI start panel).
    """
    mgr = _cowork_manager(tmp_path)
    shared = tmp_path / "shared"
    shared.mkdir()
    sid = "fresh-session"
    assert mgr.session_store.load(sid) is None  # nothing saved yet

    res = mgr.add_root(sid, str(shared), writable=False)
    assert res["ok"] is True
    paths = {r["path"] for r in res["roots"]}
    assert str(shared.resolve()) in paths

    # the grant survived: a fresh manager (no engines) still sees it, under the cowork agent
    mgr2 = _cowork_manager(tmp_path)
    persisted = {r["path"]: r for r in mgr2.get_roots(sid)}
    assert str(shared.resolve()) in persisted
    record = mgr2.session_store.load(sid)
    assert record is not None and record.agent == "cowork"


# -- Slice D: request_directory (interactive grant) ----------------------------


def test_request_directory_emits_prompt_and_returns_grant():
    captured = {}

    async def requester(args, tool_call_id=None):
        captured.update(args)
        return {"granted": True, "path": "/tmp/granted", "writable": True}

    eng = _bare_engine(directory_requester=requester)
    tc = ToolCall(
        id="c1",
        name="request_directory",
        arguments={
            "reason": "need the report",
            "path": "/tmp/granted",
            "writable": True,
        },
    )
    events = asyncio.run(_collect(eng._handle_directory_request(tc)))

    kinds = [e.type for e in events]
    assert EventType.DIRECTORY_REQUESTED in kinds
    prompt = next(e for e in events if e.type == EventType.DIRECTORY_REQUESTED)
    assert (
        prompt.data["reason"] == "need the report" and prompt.data["writable"] is True
    )
    finished = next(e for e in events if e.type == EventType.TOOL_FINISHED)
    assert finished.data["status"] == "ok"
    assert captured["reason"] == "need the report"  # the requester saw the agent's args
    # the tool result the model sees reflects the grant
    assert '"granted": true' in eng.messages[-1]["content"]


def test_request_directory_denied_returns_denied_status():
    async def requester(_args, tool_call_id=None):
        return {"granted": False, "reason": "user declined"}

    eng = _bare_engine(directory_requester=requester)
    tc = ToolCall(id="c1", name="request_directory", arguments={"reason": "x"})
    events = asyncio.run(_collect(eng._handle_directory_request(tc)))
    assert (
        next(e for e in events if e.type == EventType.TOOL_FINISHED).data["status"]
        == "denied"
    )


def test_request_directory_without_requester_is_safe_noop():
    eng = _bare_engine()  # no requester (e.g. headless)
    tc = ToolCall(id="c1", name="request_directory", arguments={"reason": "x"})
    events = asyncio.run(_collect(eng._handle_directory_request(tc)))
    # no prompt is emitted, and the tool reports it isn't available
    assert EventType.DIRECTORY_REQUESTED not in [e.type for e in events]
    assert (
        next(e for e in events if e.type == EventType.TOOL_FINISHED).data["status"]
        == "denied"
    )


# -- Universal scratch (workspace-scratch-design.md §4) -------------------------


def _repo(tmp_path, name="repo"):
    d = tmp_path / name
    d.mkdir()
    (d / "README.md").write_text("hello", encoding="utf-8")
    return d


def test_gated_session_gets_workspace_plus_scratch_roots(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path)
    sid = "sessGated1"
    engine = mgr.get_engine(sid, agent="code", workspace=str(repo))
    assert engine is not None

    roots = mgr.get_roots(sid)
    assert len(roots) == 2
    assert roots[0]["primary"] and Path(roots[0]["path"]) == repo.resolve()
    scratch = mgr.scratch_base() / sid
    assert Path(roots[1]["path"]) == scratch.resolve()
    assert roots[1]["writable"] is True
    assert scratch.is_dir()  # provisioned at engine build

    # Both roots are writable to the permission engine; elsewhere is not.
    perms = engine.permissions
    assert perms._under_writable_root(str(repo / "x.txt"))
    assert perms._under_writable_root(str(scratch / "x.txt"))
    assert not perms._under_writable_root(str(tmp_path / "elsewhere.txt"))


def test_orphan_session_keeps_scratch_primary_single_scratch_root(tmp_path):
    mgr = _cowork_manager(tmp_path)
    sid = "sessOrphan1"
    engine = mgr.get_engine(sid, agent="cowork")
    assert engine is not None
    roots = mgr.get_roots(sid)
    # ws == scratch: exactly one scratch root, never doubled.
    assert len(roots) == 1 and roots[0]["primary"] and roots[0]["writable"]
    assert Path(roots[0]["path"]) == (mgr.scratch_base() / sid).resolve()


def test_gated_session_artifacts_list_scratch_not_repo(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path)
    sid = "sessGated2"
    assert mgr.get_engine(sid, agent="code", workspace=str(repo)) is not None
    scratch = mgr.scratch_base() / sid
    (scratch / "report.html").write_text("<h1>posture</h1>", encoding="utf-8")

    arts = mgr.list_artifacts(sid)
    names = {a["name"] for a in arts}
    # The repo's files are NOT artifacts; the scratch report is.
    assert "report.html" in names and "README.md" not in names


def test_artifact_chips_resolve_in_workspace_and_scratch(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path)
    sid = "sessGated3"
    engine = mgr.get_engine(sid, agent="code", workspace=str(repo))
    assert engine is not None
    mgr.save(sid, engine)  # chips resolve against the persisted record, as after a turn
    scratch = mgr.scratch_base() / sid
    (scratch / "report.html").write_text("<h1>ok</h1>", encoding="utf-8")

    # Relative chip → workspace file (as before).
    assert mgr.read_artifact(sid, "README.md")["ok"] is True
    # Relative chip that only exists in scratch → resolves there.
    assert mgr.read_artifact(sid, "report.html")["ok"] is True
    # Absolute scratch path (what the roots context advertises) → resolves.
    assert mgr.read_artifact(sid, str(scratch / "report.html"))["ok"] is True
    # Escapes every root → refused.
    outside = tmp_path / "outside.txt"
    outside.write_text("x", encoding="utf-8")
    assert mgr.read_artifact(sid, str(outside))["ok"] is False


def test_gated_session_registers_request_directory(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path)
    engine = mgr.get_engine("sessGated4", agent="code", workspace=str(repo))
    assert engine is not None
    assert "request_directory" in engine.registry.names()


def test_gated_session_save_and_rebuild_keeps_roots_stable(tmp_path):
    """The scratch root must never persist as an 'extra' — a save/rebuild cycle used to
    re-add it as a plain folder each time (dup roots forever)."""
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path)
    sid = "sessGated5"
    engine = mgr.get_engine(sid, agent="code", workspace=str(repo))
    mgr.save(sid, engine)

    mgr2 = _cowork_manager(tmp_path)
    engine2 = mgr2.get_engine(sid)
    assert engine2 is not None
    mgr2.save(sid, engine2)
    mgr3 = _cowork_manager(tmp_path)
    roots = mgr3.get_roots(sid)
    assert [Path(r["path"]) for r in roots] == [
        repo.resolve(),
        (mgr3.scratch_base() / sid).resolve(),
    ]


# -- Root promotion (workspace-scratch-design.md §5) ----------------------------


def test_promote_workspace_repoints_orphan_session(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path, "project")
    sid = "sessPromo1"
    engine = mgr.get_engine(sid, agent="cowork")
    scratch = (mgr.scratch_base() / sid).resolve()

    res = mgr.promote_workspace(sid, str(repo))
    assert res["ok"], res
    # Live roots re-anchored: workspace first, scratch kept.
    assert [r.path for r in engine.roots][:2] == [repo.resolve(), scratch]
    assert engine.roots[0].writable and engine.roots[0].label == "workspace"
    # Persisted: the record's workspace is the promoted folder.
    rec = mgr.session_store.load(sid)
    assert Path(rec.workspace).resolve() == repo.resolve()

    # Post-turn rebuild (mark_idle evicts) → fully anchored dual-root session.
    mgr.mark_idle(sid)
    engine2 = mgr.get_engine(sid)
    assert engine2 is not engine
    roots = mgr.get_roots(sid)
    assert [Path(r["path"]) for r in roots] == [repo.resolve(), scratch]
    assert roots[0]["primary"] and roots[0]["label"] == "workspace"


def test_promote_workspace_refused_when_session_has_one(tmp_path):
    mgr = _cowork_manager(tmp_path)
    repo = _repo(tmp_path, "projectA")
    other = _repo(tmp_path, "projectB")
    sid = "sessPromo2"
    assert mgr.get_engine(sid, agent="code", workspace=str(repo)) is not None
    res = mgr.promote_workspace(sid, str(other))
    assert not res["ok"] and "already has a workspace" in res["error"]


def test_promote_workspace_is_one_way_once(tmp_path):
    mgr = _cowork_manager(tmp_path)
    a = _repo(tmp_path, "first")
    b = _repo(tmp_path, "second")
    sid = "sessPromo3"
    assert mgr.get_engine(sid, agent="cowork") is not None
    assert mgr.promote_workspace(sid, str(a))["ok"]
    res = mgr.promote_workspace(sid, str(b))
    assert not res["ok"] and "already has a workspace" in res["error"]