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()