group_fqcd_jr/tools/memory_chain_probe.py

"""记忆链路端到端探针：受理 → Worker 执行 → complete_run → 事件消费 → 抽取 → memory_unit。

为什么需要它：记忆链路上的缺陷（Worker 未注册消费、事件不携带正文而消费者期望正文、
抽取语义与幂等错误、抽取被静默跳过）在业务 Agent 接入前**不会通过现有数据自然暴露**——
`jr` 库里 `memory_unit` 一直是 0 行。探针使用生产装配（`get_agent_factory()`、真实 MySQL、
真实 `WorkerRuntime` 租约与治理链、真实 `OutboxWorker` 消费）跑完整链路并逐项断言。

判据（见 `docs/evidence/20260909-memory-baseline-before.md` 与 `-memory-chain-acceptance.md`）：
1. 运行成功后存在 `memory.extraction_requested` 事件，payload **不含正文**，只有定位字段；
2. 消费后记忆由**抽取结果**产生：受控语义键 + 结构化值，而不是用户原文整句；
3. 同一 `event_id` 重复消费不产生第二条记忆（幂等边界）。

关于装配：生产链路注入了 `IntentClassifier`，而意图分类需要已配置的模型端点；
本地 `model_endpoint_config` 为 0 行时按设计失败关闭，run 无法进入 `complete_run`。
因此探针用最小装配工厂（只注入治理）跳过意图分类，并**注入确定性替身模型服务**
（`WorkerRuntime(model_service=...)`）驱动抽取，使本探针既不依赖真实模型端点，
又能覆盖"抽取 → 记忆"这一段。生产工厂的装配完整性单独断言。

用法：
    python tools/memory_chain_probe.py            # 执行、断言并清理
    python tools/memory_chain_probe.py --keep     # 保留测试数据以便排查
"""

from __future__ import annotations

import asyncio
import sys
from uuid import uuid4

from sqlalchemy import and_, delete, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.contracts import AgentDefinition, AgentRequest, CoreResult, RequestContext
from app.core.errors import AgentTypeNotFoundError
from app.infrastructure.db import SessionFactory
from app.model.audit import InteractionAudit
from app.model.conversation import ConversationMessage
from app.model.memory import MemoryEvidence, MemoryUnit
from app.model.platform import AgentRun, DomainEventOutbox, OutboxDelivery, RequestIdempotency
from app.service.agent.base import BaseAgent
from app.service.agent.bootstrap import get_agent_factory
from app.service.agent.factory import AgentFactory
from app.service.agent.governance import AgentGovernance, PlatformGovernance
from app.service.agent_run_application_service import AgentRunApplicationService
from app.service.identity_service import IdentityService
from app.worker.memory_extraction_worker import MemoryExtractionWorker
from app.worker.runtime import WorkerRuntime

CUSTOMER_ID = 9001
AGENT_TYPE = "memory_chain_probe"
MESSAGE = "请记住：我的风险偏好是稳健型，后续建议请按稳健型说明。"
SESSION_PREFIX = "probe-"

# 抽取结果标识：P2 之后记忆键来自受控词表、内容是结构化值，不再以会话前缀命名，
# 因此清理与计数按"受控键 + 值"精确识别探针产物。
PROBE_MEMORY_KEY = "preference:risk_level"
PROBE_MEMORY_VALUE = "稳健型"
PROBE_MEMORY_TYPE = "preference"
EXTRACTION_JSON = (
    '{"memory_key": "preference:risk_level", "value": "稳健型", '
    '"memory_type": "preference", "confidence": 0.9}'
)

failures: list[str] = []


def check(condition: bool, description: str) -> None:
    print(("  PASS  " if condition else "  FAIL  ") + description)
    if not condition:
        failures.append(description)


class MemoryProbeAgent(BaseAgent):
    """最小探针 Agent：只声明意图，不调用任何工具。"""

    definition = AgentDefinition(
        agent_type=AGENT_TYPE,
        version="probe-1",
        allowed_roles=("customer",),
        allowed_portals=("api",),
        allowed_tools=(),
        supported_intents=("general",),
    )

    async def handle(self, request: AgentRequest, context: RequestContext) -> CoreResult:
        del request, context
        return CoreResult(text="已记录您的偏好说明，后续将按稳健型为您解释。")


class _StubExecution:
    def __init__(self, text: str) -> None:
        self.text = text


class StubModelService:
    """确定性替身模型：只回放抽取用的严格 JSON，不访问网络。"""

    def __init__(self, payload: str = EXTRACTION_JSON) -> None:
        self.payload = payload
        self.calls = 0

    async def generate(self, endpoints: object, prompt: str, **kwargs: object) -> _StubExecution:
        del endpoints, prompt, kwargs
        self.calls += 1
        return _StubExecution(self.payload)


class StubEndpointResolver:
    """确定性替身端点解析器：返回一个占位端点。

    `MemoryExtractionService` 在调用模型前必须先解析出可用端点，端点缺失时按设计
    失败关闭。本地 `model_endpoint_config` 为 0 行，因此探针注入本替身，使抽取路径
    可被完整验证，而不需要往库里塞一条假的模型端点配置。
    """

    def __init__(self) -> None:
        self.calls = 0

    async def resolve(self, *, agent_type: str, task_type: str) -> list[object]:
        del agent_type, task_type
        self.calls += 1
        return [object()]


def build_probe_factory() -> AgentFactory:
    """最小装配：只注入治理，避开需要模型端点的意图分类。"""
    governance: AgentGovernance = PlatformGovernance()
    factory = AgentFactory(governance=governance)
    try:
        factory.definition(AGENT_TYPE)
    except AgentTypeNotFoundError:
        factory.register(
            MemoryProbeAgent.definition,
            lambda _context: MemoryProbeAgent(MemoryProbeAgent.definition),
        )
    return factory


def _probe_memory_filter():
    return and_(
        MemoryUnit.customer_id == CUSTOMER_ID,
        MemoryUnit.memory_key == PROBE_MEMORY_KEY,
    )


async def purge_probe_residue() -> None:
    """清理历史探针残留，使本探针可重复运行。"""
    async with SessionFactory() as session, session.begin():
        memories = (await session.scalars(select(MemoryUnit).where(_probe_memory_filter()))).all()
        for memory in memories:
            await session.execute(
                delete(MemoryEvidence).where(MemoryEvidence.memory_id == memory.id)
            )
            await session.delete(memory)
        probe_runs = (
            await session.scalars(
                select(AgentRun).where(AgentRun.session_id.like(f"{SESSION_PREFIX}%"))
            )
        ).all()
        for run in probe_runs:
            events = (
                await session.scalars(
                    select(DomainEventOutbox.event_id).where(
                        DomainEventOutbox.aggregate_id == run.run_id
                    )
                )
            ).all()
            if events:
                await session.execute(
                    delete(OutboxDelivery).where(OutboxDelivery.event_id.in_(list(events)))
                )
            await session.execute(
                delete(DomainEventOutbox).where(DomainEventOutbox.aggregate_id == run.run_id)
            )
            await session.execute(
                delete(RequestIdempotency).where(RequestIdempotency.id == run.idempotency_id)
            )
            await session.execute(
                delete(ConversationMessage).where(
                    ConversationMessage.session_id == run.session_id
                )
            )
            await session.execute(
                delete(InteractionAudit).where(InteractionAudit.session_id == run.session_id)
            )
            await session.delete(run)


async def count_memories(session: AsyncSession) -> int:
    memories = (await session.scalars(select(MemoryUnit).where(_probe_memory_filter()))).all()
    return len(memories)


async def main() -> int:
    keep = "--keep" in sys.argv
    session_id = f"{SESSION_PREFIX}{uuid4()}"
    idempotency_key = f"probe{uuid4().hex}"
    context = RequestContext(
        user_id=str(CUSTOMER_ID),
        trace_id=str(uuid4()),
        roles=("customer",),
        permissions=("agent:run",),
        portal="api",
    )

    print("1) 前置检查与生产装配核对")
    await purge_probe_residue()
    check(True, "历史探针残留已清理")
    production = get_agent_factory()
    check(production._model_service is not None, "生产工厂已注入模型服务")
    check(production._tool_executor is not None, "生产工厂已注入工具执行器")
    check(production._intent_classifier is not None, "生产工厂已注入意图分类器")
    resolved = await IdentityService().resolve(context)
    if "agent:run" not in resolved.permissions:
        print("  FAIL  客户 9001 缺少 agent:run 权限；请先运行 tools/seed_test_rbac.py")
        return 1
    check(True, f"客户 {CUSTOMER_ID} 实时权限已加载（角色 {resolved.roles}）")

    factory = build_probe_factory()
    stub = StubModelService()
    resolver = StubEndpointResolver()
    check(True, f"探针工厂已装配（最小治理装配：{AGENT_TYPE}）")

    print("2) 受理运行")
    request = AgentRequest(
        agent_type=AGENT_TYPE,
        message=MESSAGE,
        session_id=session_id,
        idempotency_key=idempotency_key,
    )
    async with SessionFactory() as session:
        accepted = await AgentRunApplicationService(session, factory).accept(request, resolved)
    run_id = accepted.run_id
    check(bool(run_id), f"受理成功 run_id={run_id}")

    print("3) Worker 执行（真实租约、治理链与落库）")
    runtime = WorkerRuntime(factory=factory, model_service=stub, endpoint_resolver=resolver)
    executed = await runtime.execute(run_id)
    check(executed, "Worker 领取并执行完成")
    async with SessionFactory() as session:
        run = await session.scalar(select(AgentRun).where(AgentRun.run_id == run_id))
    check(
        run is not None and run.status == "succeeded",
        f"运行终态 = {run.status if run is not None else None}"
        + (f"（error_code={run.error_code}）" if run is not None and run.error_code else ""),
    )

    print("4) 事件契约检查")
    async with SessionFactory() as session:
        event = await session.scalar(
            select(DomainEventOutbox).where(
                DomainEventOutbox.aggregate_id == run_id,
                DomainEventOutbox.event_type == "memory.extraction_requested",
            )
        )
    check(event is not None, "complete_run 在同一事务写入了 memory.extraction_requested")
    if event is None:
        print("  说明：运行未成功时不会产生记忆事件，后续断言一并失败属预期")
        payload: dict[str, object] = {}
        event_id = run_id
    else:
        payload = dict(event.payload)
        event_id = event.event_id
    check(bool(payload) and "content" not in payload, "事件 payload 不携带正文（只带定位信息）")
    check(
        {"message_id", "customer_id"} <= set(payload),
        f"事件定位字段完整：{sorted(payload)}",
    )

    print("5) 消费事件并核对抽取结果")
    rounds = 0
    while rounds < 10:
        if not await runtime.dispatch_one(run_id=run_id):
            break
        rounds += 1
    check(rounds >= 1, f"Worker 轮询 {rounds} 轮后该 run 的事件队列清空")
    check(stub.calls >= 1, f"抽取模型被真实调用（{stub.calls} 次）")
    check(resolver.calls >= 1, f"抽取端点经解析器解析（{resolver.calls} 次）")
    async with SessionFactory() as session:
        before = await count_memories(session)
        memory = await session.scalar(select(MemoryUnit).where(_probe_memory_filter()))
    check(before >= 1, f"memory_unit 出现探针记忆（{before} 行）")
    check(
        memory is not None and memory.memory_key == PROBE_MEMORY_KEY,
        f"记忆键来自受控词表：{memory.memory_key if memory is not None else None}",
    )
    check(
        memory is not None and (memory.content or "") == PROBE_MEMORY_VALUE,
        f"记忆内容是抽取的结构化值而非用户原文：{memory.content if memory is not None else None}",
    )
    check(
        memory is not None and memory.memory_type == PROBE_MEMORY_TYPE,
        f"记忆类型与键前缀同构：{memory.memory_type if memory is not None else None}",
    )
    check(memory is not None and memory.customer_id == CUSTOMER_ID, "记忆归属为发起运行的客户")

    print("6) 幂等：同一 event_id 重复消费")
    async with SessionFactory() as session:
        # 复用同一抽取器：这样重复消费若未被幂等拦截就会真的走抽取并写入，
        # 断言才有意义（而不是因为抽取不可用而"恰好"没写）。
        worker = MemoryExtractionWorker(session, extractor=runtime.memory_extraction)
        again = await worker.handle(payload, event_id=event_id)
        await session.commit()
        after = await count_memories(session)
    check(again is False, "重复消费被幂等边界拦截")
    check(after == before, f"重复消费未新增记忆（{before} → {after}）")

    if keep:
        print(f"\n--keep 已启用：保留测试数据 session_id={session_id} run_id={run_id}")
    else:
        print("7) 清理测试数据")
        await purge_probe_residue()
        async with SessionFactory() as session:
            remaining = await count_memories(session)
        check(remaining == 0, f"探针记忆已清理（剩余 {remaining} 行）")

    print()
    if failures:
        print(f"FAILED: {len(failures)} 项未通过")
        for item in failures:
            print(f"  - {item}")
        return 1
    print("PASSED: 记忆链路端到端连通（受理 → 执行 → 事件 → 消费 → 抽取 → 记忆 → 幂等）")
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))