lambdagent/tests/test_workspace_provision.py

"""
tests/test_workspace_provision.py — 创建智能体实例时自动开辟工作目录。

布局参照 Workspace/multi-reviewer67：
  <base>/<名>/        work_dir（最终交付物 + Bash CWD）
  <base>/<名>/data    source_dir
  <base>/<名>/workspace  run_dir（bare 模式，run_* 直接建在其下）
"""
from __future__ import annotations

import json
import os

import pytest

os.environ.setdefault("AGENTPAAS_DATABASE_URL", "sqlite:///:memory:")
os.environ.setdefault("AGENTPAAS_TESTING", "1")


# ── 单元：provision 函数 ─────────────────────────────────────────────────────

def test_provision_layout(tmp_path):
    from agentpaas.engine.workspace_provision import provision_agent_workspace

    d = provision_agent_workspace("审稿助手", base=str(tmp_path))
    root = os.path.join(str(tmp_path), "审稿助手")
    assert d["work_dir"] == os.path.abspath(root)
    assert d["source_dir"] == os.path.join(os.path.abspath(root), "data")
    assert d["run_dir"] == os.path.join(os.path.abspath(root), "workspace")
    assert os.path.isdir(d["source_dir"]) and os.path.isdir(d["run_dir"])


def test_provision_duplicate_names_get_suffix(tmp_path):
    from agentpaas.engine.workspace_provision import provision_agent_workspace

    d1 = provision_agent_workspace("助手", base=str(tmp_path))
    d2 = provision_agent_workspace("助手", base=str(tmp_path))
    d3 = provision_agent_workspace("助手", base=str(tmp_path))
    assert d1["work_dir"].endswith("助手")
    assert d2["work_dir"].endswith("助手-2")
    assert d3["work_dir"].endswith("助手-3")


def test_provision_slug_sanitises_path_chars(tmp_path):
    from agentpaas.engine.workspace_provision import provision_agent_workspace

    d = provision_agent_workspace("../evil/name  v2", base=str(tmp_path))
    # 不逃逸 base；路径分隔符被清洗
    assert os.path.commonpath([d["work_dir"], str(tmp_path)]) == str(
        os.path.abspath(str(tmp_path)))
    assert "/evil/" not in d["work_dir"].replace(str(tmp_path), "")


def test_provision_failure_returns_empty(monkeypatch):
    from agentpaas.engine import workspace_provision as wp

    monkeypatch.setattr(wp.os, "makedirs",
                        lambda *a, **k: (_ for _ in ()).throw(OSError("ro")))
    d = wp.provision_agent_workspace("x", base="/nonexistent-base")
    assert d == {"work_dir": "", "source_dir": "", "run_dir": ""}


# ── 集成：两个创建端点 ───────────────────────────────────────────────────────

@pytest.fixture()
def api_client(tmp_path, monkeypatch):
    """TestClient + 内存 DB + 租户/key + workspace_base 指到 tmp。"""
    import secrets
    from fastapi.testclient import TestClient
    from agentpaas.api.app import app
    from agentpaas.api.middleware.auth import hash_key
    from agentpaas.config import settings
    from agentpaas.db.models import Database, gen_id, now_utc
    import agentpaas.db.session as _session_mod

    ws_base = tmp_path / "Workspace"
    monkeypatch.setattr(settings, "workspace_base", str(ws_base))

    prev = _session_mod._db
    _session_mod._db = Database("sqlite:///:memory:")
    db = _session_mod._db
    tid, uid = gen_id("tn_"), gen_id("usr_")
    raw_key = f"ap_{secrets.token_hex(16)}"
    now = now_utc()
    db.execute("INSERT INTO tenants (id, name, plan, status, created_at) "
               "VALUES (?, 'test', 'free', 'active', ?)", (tid, now))
    db.execute("INSERT INTO users (id, tenant_id, email, role, created_at) "
               "VALUES (?, ?, '', 'admin', ?)", (uid, tid, now))
    db.execute(
        "INSERT INTO api_keys (id, tenant_id, user_id, key_hash, key_prefix, name, "
        "scopes, rate_limit, status, created_at) "
        "VALUES (?, ?, ?, ?, ?, 'test', ?, 600, 'active', ?)",
        (gen_id("key_"), tid, uid, hash_key(raw_key), raw_key[:8],
         json.dumps(["agents:*"]), now))
    db.commit()
    with TestClient(app) as client:
        yield client, raw_key, str(ws_base), db
    _session_mod._db = prev


def _auth(k):
    return {"Authorization": f"Bearer {k}"}


_MIN_CONFIG = {"name": "目录测试", "type": "simple",
               "model": {"name": "ollama/qwen2.5:7b"}, "systemPrompt": "t"}


def test_create_agent_auto_provisions_dirs(api_client):
    client, key, ws_base, db = api_client
    r = client.post("/api/v1/agents", headers=_auth(key),
                    json={"name": "目录测试", "config": _MIN_CONFIG})
    assert r.status_code in (200, 201), r.text[:300]
    body = r.json()
    assert body["work_dir"].endswith("目录测试")
    assert body["run_dir"] == os.path.join(body["work_dir"], "workspace")
    assert os.path.isdir(body["source_dir"])

    # DB 行一致
    row = db.fetchone("SELECT work_dir, source_dir, run_dir FROM agents WHERE id=?",
                      (body["agent_id"],))
    assert row["work_dir"] == body["work_dir"]
    assert row["run_dir"] == body["run_dir"]


def test_create_agent_respects_explicit_dirs(api_client, tmp_path):
    client, key, ws_base, db = api_client
    my_dir = str(tmp_path / "my-own-dir")
    os.makedirs(my_dir)
    r = client.post("/api/v1/agents", headers=_auth(key),
                    json={"name": "显式目录", "config": _MIN_CONFIG,
                          "work_dir": my_dir})
    assert r.status_code in (200, 201), r.text[:300]
    body = r.json()
    assert body["work_dir"] == my_dir
    assert body["run_dir"] == ""  # 显式给了任一目录 → 其余不自动补
    # 没有为它建 Workspace/<名> 目录
    assert not os.path.exists(os.path.join(ws_base, "显式目录"))


# ── 模型显示与切换（模型页功能）─────────────────────────────────────────────

def test_models_in_use_and_switch(api_client):
    client, key, ws_base, db = api_client
    # 建两个不同模型的智能体
    for name, model in [("甲", {"provider": "ollama", "name": "qwen2.5:7b"}),
                        ("乙", {"provider": "dashscope", "name": "qwen-max"})]:
        cfg = {"name": name, "type": "simple", "model": model, "systemPrompt": "t"}
        r = client.post("/api/v1/agents", headers=_auth(key),
                        json={"name": name, "config": cfg})
        assert r.status_code in (200, 201), r.text[:200]

    # 一览
    r = client.get("/api/v1/providers/models-in-use", headers=_auth(key))
    assert r.status_code == 200, r.text[:200]
    items = {i["agent_name"]: i for i in r.json()["items"]}
    assert items["甲"]["provider"] == "ollama" and items["甲"]["model"] == "qwen2.5:7b"
    assert items["乙"]["provider"] == "dashscope" and items["乙"]["model"] == "qwen-max"

    # 切换甲 → dashscope/qwen-plus
    aid = items["甲"]["agent_id"]
    r = client.put(f"/api/v1/agents/{aid}/model", headers=_auth(key),
                   json={"provider": "dashscope", "name": "qwen-plus"})
    assert r.status_code == 200, r.text[:300]
    body = r.json()
    assert body["version"] == 2 and body["previous"] == "ollama/qwen2.5:7b"

    # 配置生效且其余字段保留；版本历史可回滚
    g = client.get(f"/api/v1/agents/{aid}", headers=_auth(key)).json()
    assert g["config"]["model"]["provider"] == "dashscope"
    assert g["config"]["model"]["name"] == "qwen-plus"
    assert g["config"]["systemPrompt"] == "t"
    assert g["current_version"] == 2
    rb = client.post(f"/api/v1/agents/{aid}/rollback", headers=_auth(key),
                     json={"target_version": 1})
    assert rb.status_code == 200
    g2 = client.get(f"/api/v1/agents/{aid}", headers=_auth(key)).json()
    assert g2["config"]["model"]["name"] == "qwen2.5:7b"


def test_switch_model_unknown_agent_404(api_client):
    client, key, *_ = api_client
    r = client.put("/api/v1/agents/ag_nonexistent/model", headers=_auth(key),
                   json={"provider": "ollama", "name": "x"})
    assert r.status_code == 404


# ── simple agent token 透传（P3）────────────────────────────────────────────

def test_simple_agent_usage_source_found():
    """_find_usage_source 沿包装链找到 get_usage：裸 ConversationLam、
    Guard 包装、Memory(Guard(...)) 双层包装都要命中。"""
    from lambdagent.conversation import ConversationLam
    from lambdagent.extensions import Guard, Memory

    class FakeProvider:
        default_model = "qwen-max"
        model_name = "qwen-max"
        context_window = 100000
        def chat(self, messages, **kw):
            return "ok"
        def get_usage(self):
            return {"input_tokens": 123, "output_tokens": 45,
                    "cache_creation_input_tokens": 0,
                    "cache_read_input_tokens": 0}

    lam = ConversationLam(name="t", provider=FakeProvider(), system_prompt="s")

    # 复用 _execute_agent 内的同款查找逻辑（提取为本测试的等价实现以
    # 锁定语义：_think_ref 优先 → 自身 get_usage → .agent 递归）
    def find(t, depth=0):
        if t is None or depth > 5:
            return None
        ref = getattr(t, "_think_ref", None)
        if ref is not None and callable(getattr(ref, "get_usage", None)):
            return ref
        if callable(getattr(t, "get_usage", None)):
            return t
        return find(getattr(t, "agent", None), depth + 1)

    assert find(lam) is lam
    guarded = Guard(lam, validator=lambda x: True)
    assert find(guarded) is lam
    wrapped = Memory(Guard(lam, validator=lambda x: True))
    assert find(wrapped) is lam
    assert find(wrapped).get_usage()["input_tokens"] == 123


def test_anthropic_provider_accumulates_usage():
    """AnthropicProvider 现在有 get_usage/reset_session（与 openai_compat 同形）。"""
    from lambdagent.providers.anthropic_provider import AnthropicProvider
    from lambdagent.providers.base import ProviderConfig

    p = AnthropicProvider(ProviderConfig(model="claude-sonnet-4-6"))
    u = p.get_usage()
    assert u == {"input_tokens": 0, "output_tokens": 0,
                 "cache_creation_input_tokens": 0,
                 "cache_read_input_tokens": 0}
    p._usage_input += 10; p._usage_output += 5
    assert p.get_usage()["input_tokens"] == 10
    p.reset_session()
    assert p.get_usage()["output_tokens"] == 0