mnote/wolai-backend/tests/test_ai_document_agent.py

import contextlib
import unittest
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

from app.services.ai_document_agent import (
    DocumentAiEmitter,
    DocumentAiContextPayload,
    DocumentAiRunContext,
    DocumentAiMessage,
    DocumentAiRunRequest,
    LocalBridgeRuntime,
    build_agent_usage_from_response_usage,
    build_document_ai_config_payload,
    build_document_model_provider,
    build_document_agent_instructions,
    build_forward_headers,
    run_document_agent_stream,
    resolve_document_agent_request_model,
    resolve_document_agent_model,
)


class AiDocumentAgentTests(unittest.TestCase):
    def test_build_document_agent_instructions_includes_core_context(self):
        request = DocumentAiRunRequest(
            userId="user-1",
            messages=[DocumentAiMessage(role="user", content="帮我改写第一段")],
            context=DocumentAiContextPayload(
                documentId="doc-1",
                documentBlocks=[
                    {
                        "id": "block-1",
                        "type": "paragraph",
                        "content": [{"text": "第一段正文"}],
                        "children": [],
                    }
                ],
                pageOptions={"wideLayout": True},
                outline={"items": ["第一段"]},
            ),
        )

        instructions = build_document_agent_instructions(request)

        self.assertIn("documentId=doc-1", instructions)
        self.assertIn("documentBlockSummaries=", instructions)
        self.assertIn("第一段正文", instructions)
        self.assertIn("pageOptions=", instructions)
        self.assertIn("outline=", instructions)

    def test_build_document_agent_instructions_includes_profile_policy(self):
        request = DocumentAiRunRequest(
            userId="user-1",
            profileId="page_writer_ai_first",
            messages=[DocumentAiMessage(role="user", content="帮我补一段总结")],
            context=DocumentAiContextPayload(
                documentId="doc-1",
                documentBlocks=[],
            ),
        )

        instructions = build_document_agent_instructions(request)

        self.assertIn("AI 优先直接进入主编辑区起草", instructions)
        self.assertIn("documentId=doc-1", instructions)

    def test_build_document_agent_instructions_includes_title_write_hard_rules(self):
        request = DocumentAiRunRequest(
            userId="user-1",
            messages=[DocumentAiMessage(role="user", content="请把当前页面标题改成“新标题”。")],
            context=DocumentAiContextPayload(
                documentId="doc-1",
                documentBlocks=[],
            ),
        )

        instructions = build_document_agent_instructions(request)

        self.assertIn("你必须调用一次 slash_run", instructions)
        self.assertIn("/rename doc-1 <新标题>", instructions)

    def test_build_forward_headers_only_keeps_allowed_headers(self):
        forwarded = build_forward_headers(
            {
                "authorization": "Bearer token",
                "cookie": "a=1",
                "x-request-id": "req-1",
                "x-trace-id": "trace-1",
                "x-mnote-source-channel": "next_route",
                "user-agent": "Vitest",
                "x-ignore-me": "nope",
            }
        )

        self.assertEqual(forwarded["authorization"], "Bearer token")
        self.assertEqual(forwarded["cookie"], "a=1")
        self.assertEqual(forwarded["x-request-id"], "req-1")
        self.assertNotIn("x-ignore-me", forwarded)

    def test_resolve_document_agent_model_uses_prefixed_request_model_first(self):
        resolved = resolve_document_agent_model(
            request_model="gh/gpt-4.1",
            configured_default_model="ju/gpt-5.4",
            env={},
            available_models_loader=lambda _base_url, _api_key: [],
        )

        self.assertEqual(resolved, "gh/gpt-4.1")

    def test_resolve_document_agent_request_model_prefers_model_key_semantic_name(self):
        resolved = resolve_document_agent_request_model(
            request_model=None,
            request_model_key="gpt-5.3-codex",
            configured_default_model="ju/gpt-5.4",
            env={
                "OPENAI_BASE_URL": "http://127.0.0.1:20128/v1",
                "OPENAI_API_KEY": "sk-test",
            },
            available_models_loader=lambda _base_url, _api_key: [
                "slow",
                "fast",
                "codex",
                "ju/gpt-5.4",
            ],
        )

        self.assertEqual(resolved, "gpt-5.3-codex")

    def test_resolve_document_agent_model_prefers_responses_provider_for_bare_default(self):
        resolved = resolve_document_agent_model(
            request_model=None,
            configured_default_model="",
            env={
                "CLAW_DEFAULT_MODEL": "gpt-5.4",
                "OPENAI_BASE_URL": "http://127.0.0.1:20128/v1",
                "OPENAI_API_KEY": "sk-test",
            },
            available_models_loader=lambda _base_url, _api_key: [
                "gh/gpt-5.4",
                "ju/gpt-5.4",
                "gmn/gpt-5.4",
            ],
        )

        self.assertEqual(resolved, "ju/gpt-5.4")

    def test_resolve_document_agent_model_requires_explicit_prefix_when_bare_model_still_ambiguous(self):
        with self.assertRaisesRegex(RuntimeError, "provider/model"):
            resolve_document_agent_model(
                request_model=None,
                configured_default_model="",
                env={
                    "CLAW_DEFAULT_MODEL": "gpt-5.4",
                    "OPENAI_BASE_URL": "http://127.0.0.1:20128/v1",
                    "OPENAI_API_KEY": "sk-test",
                },
                available_models_loader=lambda _base_url, _api_key: [
                    "gh/gpt-5.4",
                    "gmn/gpt-5.4",
                ],
            )

    def test_resolve_document_agent_model_fails_when_no_model_is_configured(self):
        with self.assertRaisesRegex(RuntimeError, "未配置文档页 AI model"):
            resolve_document_agent_model(
                request_model=None,
                configured_default_model="",
                env={},
                available_models_loader=lambda _base_url, _api_key: [],
            )

    def test_build_document_model_provider_supports_gateway_prefixed_models(self):
        provider = build_document_model_provider(
            request_model="ju/gpt-5.4",
            configured_default_model="",
            env={
                "OPENAI_BASE_URL": "http://127.0.0.1:20128/v1",
                "OPENAI_API_KEY": "sk-test",
            },
            available_models_loader=lambda _base_url, _api_key: [
                "ju/gpt-5.4",
                "gh/gpt-4.1",
            ],
        )

        ju_model = provider.get_model("ju/gpt-5.4")
        gh_model = provider.get_model("gh/gpt-4.1")

        self.assertEqual(getattr(ju_model, "model", None), "ju/gpt-5.4")
        self.assertEqual(getattr(gh_model, "model", None), "gh/gpt-4.1")

    def test_build_agent_usage_from_response_usage_falls_back_to_input_plus_output(self):
        class FakeUsage:
            input_tokens = 12
            output_tokens = 5
            total_tokens = None
            input_tokens_details = None
            output_tokens_details = None

        usage = build_agent_usage_from_response_usage(FakeUsage())

        self.assertEqual(usage.input_tokens, 12)
        self.assertEqual(usage.output_tokens, 5)
        self.assertEqual(usage.total_tokens, 17)

    def test_build_agent_usage_from_response_usage_ignores_invalid_total_tokens(self):
        class FakeUsage:
            input_tokens = 8
            output_tokens = 3
            total_tokens = "not-a-number"
            input_tokens_details = None
            output_tokens_details = None

        usage = build_agent_usage_from_response_usage(FakeUsage())

        self.assertEqual(usage.input_tokens, 8)
        self.assertEqual(usage.output_tokens, 3)
        self.assertEqual(usage.total_tokens, 11)

    def test_build_document_ai_config_payload_exposes_model_keys_profiles_and_tools(self):
        payload = build_document_ai_config_payload(
            env={
                "OPENAI_BASE_URL": "http://127.0.0.1:20128/v1",
                "CLAW_DEFAULT_MODEL": "gpt-5.4",
            }
        )

        self.assertEqual(payload["provider"], "online")
        self.assertEqual(payload["baseUrl"], "http://localhost:20128/v1")
        self.assertEqual(payload["defaultModelKey"], "gpt-5.4")
        self.assertEqual(payload["defaultProfileId"], "page_writer_ai_first")
        self.assertTrue(payload["sessionEnabled"])
        self.assertTrue(any(item["key"] == "gpt-5.4" for item in payload["models"]))
        self.assertTrue(
            any(
                item["key"] == "gpt-5.4"
                and item["resolvedCombo"] == "slow"
                and item["resolvedRuntimeModel"] == "gpt-5.4"
                for item in payload["models"]
            )
        )
        self.assertTrue(any(item["id"] == "page_writer_ai_first" for item in payload["profiles"]))
        self.assertTrue(any(item["name"] == "doc_get" for item in payload["tools"]))


class AiDocumentAgentBridgeRuntimeTests(unittest.IsolatedAsyncioTestCase):
    async def test_local_bridge_runtime_execute_tool_uses_local_runtime_result_mode(self):
        bridge = LocalBridgeRuntime()
        ctx = DocumentAiRunContext(
            user_id="user-1",
            request_id="req-1",
            trace_id="trace-1",
            session_id=None,
            workspace_id="ws-1",
            document_id="doc-1",
            document_blocks=[],
            node=None,
            subtree=None,
            outline=None,
            evidence=None,
            page_options=None,
            editor_runtime_page_options=None,
            forward_headers={},
            emitter=DocumentAiEmitter(),
            bridge=bridge,
        )

        mock_run_bridge_runtime = AsyncMock(
            return_value={
                "ok": True,
                "result": {
                    "ok": True,
                    "parsed": {
                        "command": "rename_doc",
                    },
                },
            }
        )
        with patch(
            "app.services.ai_document_agent.run_bridge_runtime",
            new=mock_run_bridge_runtime,
        ):
            result = await bridge.execute_tool(
                tool="slash_run",
                invocation_kind="command",
                context=ctx,
                args_json={"text": "/rename doc-1 新标题"},
                data={"source": "ai-orchestrator"},
                target={"pageId": "doc-1"},
            )

        self.assertTrue(result["ok"])
        self.assertEqual(result["parsed"]["command"], "rename_doc")
        payload = mock_run_bridge_runtime.await_args.args[0]
        self.assertEqual(payload["kind"], "tool")
        self.assertEqual(payload["tool"]["tool"], "slash_run")
        self.assertEqual(payload["tool"]["kind"], "command")
        self.assertEqual(payload["tool"]["mode"], "result")
        self.assertEqual(payload["tool"]["target"]["pageId"], "doc-1")
        self.assertEqual(payload["context"]["source"]["channel"], "ai_orchestrator")
        self.assertEqual(payload["context"]["source"]["client"], "wolai-backend")

    async def test_run_document_agent_stream_does_not_require_legacy_mnote_web_config(self):
        request = DocumentAiRunRequest(
            userId="user-1",
            maxSteps=1,
            messages=[DocumentAiMessage(role="user", content="请总结当前页面")],
            context=DocumentAiContextPayload(
                documentId="doc-1",
                documentBlocks=[],
            ),
        )

        class FakeRunner:
            @staticmethod
            async def run(*_args, **_kwargs):
                return SimpleNamespace(final_output="已完成")

        with (
            patch("app.services.ai_document_agent.sdk_available", return_value=True),
            patch(
                "app.services.ai_document_agent.has_configured_openai_api_key",
                return_value=True,
            ),
            patch(
                "app.services.ai_document_agent.build_document_agent",
                return_value=object(),
            ),
            patch(
                "app.services.ai_document_agent.build_document_model_provider",
                return_value=object(),
            ),
            patch(
                "app.services.ai_document_agent.build_session",
                return_value=None,
            ),
            patch(
                "app.services.ai_document_agent.RunConfig",
                side_effect=lambda **kwargs: kwargs,
            ),
            patch(
                "app.services.ai_document_agent.Runner",
                new=FakeRunner,
            ),
            patch(
                "app.services.ai_document_agent.trace",
                new=lambda *_args, **_kwargs: contextlib.nullcontext(),
            ),
        ):
            chunks: list[str] = []
            async for chunk in run_document_agent_stream(request, source_headers={}):
                chunks.append(chunk)

        text = "".join(chunks)
        self.assertIn("event: assistant_message", text)
        self.assertIn("event: completion", text)
        self.assertIn("已完成", text)


if __name__ == "__main__":
    unittest.main()