feat: add agent session usage status reporting

Track per-session model and token usage so users can inspect context pressure and cumulative usage with /session_status.
2026-06-07 08:41:16 +08:00 · 2026-04-26 08:19:05 +08:00
parent 79bfeaf2af
commit 0277288a41
5 changed files with 552 additions and 12 deletions
--- a/tests/test_agent_session_status.py
+++ b/tests/test_agent_session_status.py
@@ -0,0 +1,106 @@
+import asyncio
+import unittest
+from datetime import datetime
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from langchain.agents.middleware.types import ModelRequest, ModelResponse
+from langchain_core.messages import AIMessage
+
+from app.agent.middleware.usage import UsageMiddleware
+from app.chain.message import MessageChain
+from app.schemas.types import MessageChannel
+
+
+class TestAgentSessionStatus(unittest.TestCase):
+    def test_usage_middleware_records_usage_metadata(self):
+        snapshots = []
+        middleware = UsageMiddleware(on_usage=snapshots.append)
+        request = ModelRequest(
+            model=SimpleNamespace(
+                model="gpt-4o-mini", profile={"max_input_tokens": 128000}
+            ),
+            messages=[],
+            state={},
+            runtime=None,
+        )
+        response = ModelResponse(
+            result=[
+                AIMessage(
+                    content="ok",
+                    usage_metadata={
+                        "input_tokens": 1200,
+                        "output_tokens": 300,
+                        "total_tokens": 1500,
+                    },
+                )
+            ]
+        )
+
+        async def handler(_: ModelRequest):
+            return response
+
+        result = asyncio.run(middleware.awrap_model_call(request, handler))
+
+        self.assertIs(result, response)
+        self.assertEqual(len(snapshots), 1)
+        self.assertEqual(snapshots[0]["model"], "gpt-4o-mini")
+        self.assertEqual(snapshots[0]["context_window_tokens"], 128000)
+        self.assertEqual(snapshots[0]["input_tokens"], 1200)
+        self.assertEqual(snapshots[0]["output_tokens"], 300)
+        self.assertEqual(snapshots[0]["total_tokens"], 1500)
+        self.assertAlmostEqual(snapshots[0]["context_usage_ratio"], 1200 / 128000)
+
+    def test_remote_session_status_sends_usage_summary(self):
+        chain = MessageChain()
+        chain._user_sessions["10001"] = ("session-1", datetime.now())
+        status = {
+            "session_id": "session-1",
+            "model": "gpt-4o-mini",
+            "context_window_tokens": 128000,
+            "last_input_tokens": 1200,
+            "last_output_tokens": 300,
+            "last_total_tokens": 1500,
+            "last_context_usage_ratio": 1200 / 128000,
+            "total_input_tokens": 4500,
+            "total_output_tokens": 1500,
+            "total_tokens": 6000,
+            "model_call_count": 4,
+            "last_updated_at": "2026-04-26 12:34:56",
+            "is_processing": True,
+            "pending_messages": 2,
+        }
+
+        with (
+            patch(
+                "app.chain.message.agent_manager.get_session_status",
+                return_value=status,
+            ),
+            patch.object(chain, "post_message") as post_message,
+        ):
+            chain.remote_session_status(
+                channel=MessageChannel.Telegram,
+                userid="10001",
+                source="telegram-test",
+            )
+
+        notification = post_message.call_args.args[0]
+        self.assertEqual(notification.title, "当前智能体会话状态")
+        self.assertIn("session-1", notification.text)
+        self.assertIn("gpt-4o-mini", notification.text)
+        self.assertIn("1,200 / 128,000 (0.94%)", notification.text)
+        self.assertIn("输入 4,500 / 输出 1,500 / 总计 6,000", notification.text)
+        self.assertIn("运行中", notification.text)
+
+    def test_remote_session_status_handles_missing_session(self):
+        chain = MessageChain()
+
+        with patch.object(chain, "post_message") as post_message:
+            chain.remote_session_status(
+                channel=MessageChannel.Telegram,
+                userid="10001",
+                source="telegram-test",
+            )
+
+        notification = post_message.call_args.args[0]
+        self.assertEqual(notification.title, "您当前没有活跃的智能体会话")