Artemis-Eapil/artemis/plugins/evidence.py

"""Task-local receipts and verified attachments; traces remain the execution ledger."""

import base64
import hashlib
import json
from pathlib import Path
import time
from typing import Any

from artemis.plugins.models import EvidenceSource, PluginCallResult

MAX_ATTACHMENT_BYTES = 64 * 1024 * 1024


def source_identity(plugin_id: str, source: EvidenceSource | None) -> str:
    if source is None:
        return f"plugin:{plugin_id}"
    identity = source.model_dump(exclude={"label", "captured_at"})
    digest = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()[:16]
    return f"plugin:{plugin_id}:{digest}"


def validate_tools(manifest, tools) -> None:
    """V2 capability declarations must describe the actual MCP surface."""
    capabilities = manifest.capabilities
    if capabilities is None:
        return
    actual = {tool.name: tool for tool in tools}
    missing = (set(capabilities.tools) | set(capabilities.required_tools)) - actual.keys()
    undeclared = actual.keys() - capabilities.tools.keys()
    if missing or undeclared:
        raise ValueError(
            f"插件工具与清单不一致；缺少：{', '.join(sorted(missing)) or '无'}；"
            f"未声明：{', '.join(sorted(undeclared)) or '无'}"
        )
    for name, spec in capabilities.tools.items():
        tool = actual[name]
        if spec.target_argument and spec.target_argument not in tool.inputSchema.get(
            "properties", {}
        ):
            raise ValueError(f"工具 {name} 缺少目标参数 {spec.target_argument}")
        if spec.observation and not (tool.annotations and tool.annotations.readOnlyHint):
            raise ValueError(f"观察工具 {name} 的 MCP 只读声明与清单不一致")


def decode_receipt(outcome: PluginCallResult) -> dict[str, Any]:
    raw = (outcome.structured_content or {}).get("artemis", {})
    if not isinstance(raw, dict):
        raise ValueError("插件证据回执必须是对象")
    if not raw:
        return {}
    if raw.get("version") != 1:
        raise ValueError("不支持的插件证据回执版本")
    if raw.get("source"):
        outcome.source = EvidenceSource.model_validate(raw["source"])
    if raw.get("execution") in {"not_started", "completed", "partial", "unknown"}:
        outcome.execution = raw["execution"]
    if raw.get("side_effect") in {"none", "possible", "confirmed"}:
        outcome.side_effect = raw["side_effect"]
    if outcome.status != "success":
        outcome.error_code = str(raw.get("error_code") or "tool_error")[:100]
    # A plugin cannot authorize replay of a possibly completed write.
    outcome.retryable = bool(raw.get("retryable")) and outcome.execution == "not_started"
    outcome.requires_observation = outcome.execution in {"partial", "unknown"}
    return raw


def persist_attachments(root: Path, plugin_id: str, receipts: list) -> list[dict]:
    directory = root / "artifacts" / "plugins" / plugin_id
    directory.mkdir(parents=True, exist_ok=True)
    if directory.resolve() != directory.absolute() or not directory.resolve().is_relative_to(
        root.resolve()
    ):
        raise ValueError("附件目录不能通过链接指向任务目录之外")
    verified = []
    for item in receipts:
        if not isinstance(item, dict):
            raise ValueError("附件回执必须是对象")
        name = str(item.get("name") or "")
        if (
            not name
            or len(name) > 180
            or name != Path(name).name
            or any(c in name for c in '/\\:*?"<>|')
            or name in {".", ".."}
            or name.rstrip(" .") != name
            or name.split(".")[0].upper()
            in {
                "CON",
                "PRN",
                "AUX",
                "NUL",
                *(f"COM{i}" for i in range(10)),
                *(f"LPT{i}" for i in range(10)),
            }
            or any(ord(c) < 32 for c in name)
        ):
            raise ValueError("附件必须使用普通文件名")
        path = directory / name
        if path.is_symlink() or path.resolve().parent != directory.resolve():
            raise ValueError("附件不能指向任务目录之外")
        if "data" in item:
            encoded = item["data"]
            if not isinstance(encoded, str) or len(encoded) > MAX_ATTACHMENT_BYTES * 4 // 3 + 4:
                raise ValueError("附件超过大小限制")
            data = base64.b64decode(encoded, validate=True)
            if len(data) > MAX_ATTACHMENT_BYTES:
                raise ValueError("附件超过大小限制")
            if hashlib.sha256(data).hexdigest() != item.get("sha256"):
                raise ValueError("插件附件读回校验失败")
            if path.exists():
                if path.read_bytes() != data:
                    raise ValueError("附件重名，不能覆盖已有证据")
            else:
                with path.open("xb") as handle:
                    handle.write(data)
        if not path.is_file() or path.stat().st_size > MAX_ATTACHMENT_BYTES:
            raise ValueError("插件附件不存在或超过大小限制")
        digest = hashlib.sha256(path.read_bytes()).hexdigest()
        if item.get("sha256") != digest:
            raise ValueError("插件附件读回校验失败")
        verified.append(
            {
                "path": path.relative_to(root).as_posix(),
                "name": name,
                "bytes": path.stat().st_size,
                "sha256": digest,
                "media_type": str(item.get("media_type") or "application/octet-stream"),
            }
        )
    return verified


def persist_receipt(ctx, outcome: PluginCallResult, arguments: dict, started_at: float) -> None:
    engine = getattr(ctx, "data_engine", None)
    if not engine or not isinstance(getattr(engine, "base_dir", None), (str, Path)):
        return
    root = Path(engine.base_dir)
    directory = root / "evidence"
    directory.mkdir(parents=True, exist_ok=True)
    receipt = outcome.model_dump(exclude={"content"})
    # The complete redacted result is already in the tool trace. Keep the index small.
    receipt.pop("structured_content", None)
    receipt.update({"arguments": arguments, "started_at": started_at, "ended_at": time.time()})
    with (directory / "calls.jsonl").open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(receipt, ensure_ascii=False) + "\n")