script-build-and-framework-tmp/script_build_host/tests/test_retrieval_adapters.py

from __future__ import annotations

import json
from hashlib import sha256
from types import SimpleNamespace

import httpx
import pytest

from script_build_host.adapters.retrieval import (
    ExternalRetrievalAdapter,
    FileDecodeRetrievalAdapter,
    FileKnowledgeRetrievalAdapter,
    HttpDecodeRetrievalAdapter,
    PatternRetrievalAdapter,
    SafeHttpClient,
    SafeImageAdapter,
)
from script_build_host.adapters.uploaded_topic import SqlUploadedTopicGateway
from script_build_host.domain.errors import ProtocolViolation, UnsafeOutboundTarget
from script_build_host.infrastructure.outbound import OutboundPolicy
from script_build_host.repositories.legacy_input import LegacySqlAlchemyInputReader


async def _public_resolver(_: str, __: int) -> tuple[str, ...]:
    return ("93.184.216.34",)


class _RawArtifacts:
    def __init__(self) -> None:
        self.values: list[tuple[str, bytes]] = []

    async def freeze_bytes(self, content: bytes, *, media_type: str) -> str:
        self.values.append((media_type, content))
        return "script-build://raw-artifacts/sha256/" + sha256(content).hexdigest()


@pytest.mark.asyncio
async def test_http_retrieval_revalidates_redirect_and_freezes_response_metadata() -> None:
    calls: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(str(request.url))
        return httpx.Response(
            200,
            json={"data": [{"channel_content_id": "post-1", "title": "case"}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await ExternalRetrievalAdapter(
            "https://api.example.com/search", safe, raw_store
        ).retrieve(query={"keyword": "case"}, snapshot=SimpleNamespace())
    assert calls == ["https://api.example.com/search"]
    assert result.source_refs == ("external:post-1",)
    assert str(result.metadata["response_sha256"]).startswith("sha256:")
    assert result.raw_artifact_ref is not None
    assert result.raw_artifact_ref.endswith(
        str(result.metadata["response_sha256"]).removeprefix("sha256:")
    )
    assert raw_store.values[0][0] == "application/json"

    def redirect(_: httpx.Request) -> httpx.Response:
        return httpx.Response(302, headers={"location": "https://127.0.0.1/private"})

    async with httpx.AsyncClient(transport=httpx.MockTransport(redirect)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        with pytest.raises(UnsafeOutboundTarget):
            await safe.request("GET", "https://api.example.com/start")

    def oversized(_: httpx.Request) -> httpx.Response:
        return httpx.Response(200, content=b"123456")

    async with httpx.AsyncClient(transport=httpx.MockTransport(oversized)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
            max_response_bytes=5,
        )
        with pytest.raises(ProtocolViolation, match="byte limit"):
            await safe.request("GET", "https://api.example.com/large")


@pytest.mark.asyncio
async def test_external_response_secrets_are_redacted_before_evidence_summary() -> None:
    def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "data": [
                    {
                        "id": "p1",
                        "authorization": "Bearer raw-token",
                        "url": "https://source.example/item?token=secret",
                    }
                ]
            },
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await ExternalRetrievalAdapter(
            "https://api.example.com/search", safe, raw_store
        ).retrieve(query={"keyword": "case"}, snapshot=SimpleNamespace())
    assert "raw-token" not in result.summary
    assert "secret" not in result.summary
    frozen = raw_store.values[0][1].decode("utf-8")
    assert "raw-token" not in frozen
    assert "secret" not in frozen


@pytest.mark.asyncio
async def test_external_large_response_keeps_bounded_summary_and_complete_raw_json() -> None:
    rows = [{"id": f"p{index}", "text": "x" * 2_000} for index in range(20)]

    def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, json={"data": rows})

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await ExternalRetrievalAdapter(
            "https://api.example.com/search", safe, raw_store
        ).retrieve(query={"keyword": "case"}, snapshot=SimpleNamespace())

    assert len(result.summary) <= 20_000
    assert json.loads(result.summary)["truncated_for_context"] is True
    assert len(json.loads(raw_store.values[0][1])["data"]) == 20


@pytest.mark.asyncio
async def test_pattern_polling_and_image_limits_use_mock_transport() -> None:
    poll_count = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal poll_count
        if request.url.path == "/pattern":
            return httpx.Response(
                200,
                json={
                    "status": "pending",
                    "session_id": "task-1",
                    "poll_url": "https://api.example.com/poll/task-1",
                },
            )
        if request.url.path.startswith("/poll"):
            poll_count += 1
            return httpx.Response(200, json={"status": "done", "results": [{"id": 2}]})
        return httpx.Response(
            200,
            content=b"\x89PNG\r\n\x1a\nimage",
            headers={"content-type": "image/png"},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await PatternRetrievalAdapter(
            "https://api.example.com/pattern",
            safe,
            raw_store,
            poll_interval_seconds=0.001,
        ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())
        images = await SafeImageAdapter(
            safe,
            raw_store,
            max_images=1,
            max_image_bytes=20,
            max_total_image_bytes=20,
        ).load(urls=["https://api.example.com/image"], snapshot=SimpleNamespace())
    assert poll_count == 1
    assert result.source_refs == ("pattern:2",)
    assert result.raw_artifact_ref is not None
    assert images[0]["digest"].startswith("sha256:")
    assert "url" not in images[0]
    assert images[0]["raw_artifact_ref"].startswith("script-build://raw-artifacts/")
    assert [media_type for media_type, _ in raw_store.values] == [
        "application/json",
        "image/png",
    ]


@pytest.mark.asyncio
async def test_pattern_timeout_is_frozen_as_a_limitation() -> None:
    def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, json={"status": "pending", "session_id": "task-1"})

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        result = await PatternRetrievalAdapter(
            "https://api.example.com/pattern",
            safe,
            _RawArtifacts(),
            poll_interval_seconds=0.001,
            poll_timeout_seconds=0,
        ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())
    assert result.source_refs == ()
    assert result.limitations == ("timeout",)
    assert result.metadata["task_id"] == "task-1"


@pytest.mark.asyncio
async def test_pattern_result_keeps_five_candidates_and_freezes_all_items() -> None:
    candidates = [
        {
            "itemset": [{"dimension": f"dimension-{index}"}],
            "target_items": [{"dimension": f"dimension-{index}"}],
            "support": 10 - index,
        }
        for index in range(6)
    ]

    async def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json=[{"status": "success", "data": {"candidates": candidates}}],
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await PatternRetrievalAdapter(
            "https://api.example.com/pattern", safe, raw_store
        ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())

    payload = json.loads(result.summary)[0]
    assert payload["data"]["candidates_total"] == 6
    assert payload["data"]["candidates_returned"] == 5
    assert len(payload["data"]["candidates"]) == 5
    assert all("target_items" not in item for item in payload["data"]["candidates"])
    frozen = json.loads(raw_store.values[0][1])
    assert len(frozen[0]["data"]["candidates"]) == 6
    assert result.raw_artifact_ref is not None


@pytest.mark.asyncio
async def test_pattern_result_has_a_hard_context_size_limit() -> None:
    candidate = {
        "itemset": [{"dimension": "x" * 10_000}],
        "target_items": [{"dimension": "x" * 10_000}],
        "examples": ["post-1"],
    }

    async def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json=[
                {
                    "status": "success",
                    "output": "large result",
                    "data": {"candidates_total": 193, "candidates": [candidate] * 20},
                }
            ],
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await PatternRetrievalAdapter(
            "https://api.example.com/pattern", safe, raw_store
        ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())

    assert len(result.summary) <= 20_100
    payload = json.loads(result.summary)[0]
    assert payload["data"]["candidates_total"] == 193
    assert payload["data"]["truncated_for_context"] is True
    assert len(raw_store.values[0][1]) > len(result.summary.encode("utf-8"))


@pytest.mark.asyncio
async def test_file_decode_and_knowledge_freeze_file_hash_and_rank(tmp_path) -> None:
    decode_path = tmp_path / "decode.json"
    decode_path.write_text(
        json.dumps({"items": [{"post_id": 4, "account": "acct", "text": "focus"}]}),
        encoding="utf-8",
    )
    knowledge_path = tmp_path / "knowledge.json"
    knowledge_path.write_text(json.dumps([{"id": "k1", "title": "focus method"}]), encoding="utf-8")
    decode_adapter = FileDecodeRetrievalAdapter(decode_path)
    snapshot = SimpleNamespace(
        model_manifest={"embedding_model": "frozen-embed"},
        datasource_manifest={},
    )
    decode = await decode_adapter.retrieve(
        query={
            "keyword": "focus",
            "account_name": "acct",
            "top_k": 3,
            "return_field": "主脉络",
        },
        snapshot=snapshot,
    )
    knowledge = await FileKnowledgeRetrievalAdapter(knowledge_path).retrieve(
        query={"keyword": "focus", "max_count": 3}, snapshot=snapshot
    )
    assert decode.metadata["embedding_model"] == "frozen-embed"
    assert decode.metadata["ranks"] == [1]
    assert decode.metadata["scores"] == [1.0]
    assert str(knowledge.metadata["file_sha256"]).startswith("sha256:")

    frozen_digest = decode.metadata["index_sha256"]
    decode_path.write_text(json.dumps({"items": []}), encoding="utf-8")
    replay = await decode_adapter.retrieve(
        query={"keyword": "focus", "top_k": 3, "return_field": "主脉络"},
        snapshot=SimpleNamespace(
            model_manifest={},
            datasource_manifest={"decode_index": {"sha256": frozen_digest}},
        ),
    )
    assert replay.source_refs == ("decode:4",)
    with pytest.raises(ProtocolViolation, match="changed"):
        await FileDecodeRetrievalAdapter(decode_path).retrieve(
            query={"keyword": "focus", "top_k": 3, "return_field": "主脉络"},
            snapshot=SimpleNamespace(
                model_manifest={},
                datasource_manifest={"decode_index": {"sha256": frozen_digest}},
            ),
        )


@pytest.mark.asyncio
async def test_http_decode_preserves_legacy_query_and_ranked_scores() -> None:
    requests: list[dict[str, object]] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(json.loads(request.content))
        return httpx.Response(
            200,
            json={"items": [{"post_id": "p1", "score": 0.91, "text": "safe"}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        safe = SafeHttpClient(
            client,
            OutboundPolicy(frozenset({"decode.example"}), resolver=_public_resolver),
        )
        raw_store = _RawArtifacts()
        result = await HttpDecodeRetrievalAdapter(
            "https://decode.example/search",
            safe,
            raw_store,
        ).retrieve(
            query={
                "return_field": "主脉络",
                "keyword": "focus",
                "account_name": "acct",
                "match_fields": {},
                "top_k": 3,
            },
            snapshot=SimpleNamespace(
                datasource_manifest={"decode_index": {"sha256": "sha256:" + "1" * 64}},
                model_manifest={"embedding_model": "embed-v1"},
            ),
        )
    assert requests[0]["top_k"] == 3
    assert result.source_refs == ("decode:p1",)
    assert result.metadata["scores"] == [0.91]
    assert result.metadata["ranks"] == [1]
    assert json.loads(raw_store.values[0][1])["items"][0]["post_id"] == "p1"
    assert result.raw_artifact_ref is not None


@pytest.mark.asyncio
async def test_uploaded_topic_parser_keeps_old_shape_without_writing() -> None:
    gateway = SqlUploadedTopicGateway(SimpleNamespace())  # type: ignore[arg-type]
    parsed = await gateway.parse(
        {
            "选题融合": "topic",
            "target_post": {"channel_account_name": "acct"},
            "灵感点": [
                {
                    "点": "point",
                    "实质": {"具体元素": [{"名称": "element", "说明": "note"}]},
                }
            ],
        }
    )
    assert parsed["account_name"] == "acct"
    assert parsed["item_count"] == 1
    assert parsed["points"][0]["items"][0]["dimension"] == "实质"


@pytest.mark.asyncio
async def test_uploaded_topic_create_writes_a_readable_legacy_graph(database) -> None:
    _, sessions = database
    gateway = SqlUploadedTopicGateway(sessions)
    parsed = await gateway.parse(
        {
            "选题融合": "uploaded topic",
            "target_post": {"channel_account_name": "acct"},
            "关键点": [
                {
                    "点": "point",
                    "形式": [{"名称": "element", "说明": "note"}],
                }
            ],
        }
    )
    created = await gateway.create(parsed, account_name=None)
    graph = await LegacySqlAlchemyInputReader(sessions).read_topic_graph(
        execution_id=0,
        topic_build_id=created["topic_build_id"],
        topic_id=created["topic_id"],
    )
    assert graph["execution"] == {"id": 0, "status": "upload"}
    assert graph["topic"]["result"] == "uploaded topic"
    assert graph["points"][0]["item_ids"] == [graph["composition_items"][0]["id"]]