From d711eefe966b8cdce264be62cf30d614ae48b960 Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Fri, 3 Apr 2026 18:03:31 +0100
Subject: [PATCH] patch: agent usage limits

---
 application/api/answer/routes/answer.py | 4 ++++
 application/api/answer/routes/stream.py | 8 ++++++++
 tests/agents/test_mcp_tool.py           | 1 +
 3 files changed, 13 insertions(+)

diff --git a/application/api/answer/routes/answer.py b/application/api/answer/routes/answer.py
index 5fa7199f..f2ec9052 100644
--- a/application/api/answer/routes/answer.py
+++ b/application/api/answer/routes/answer.py
@@ -85,6 +85,10 @@ class AnswerResource(Resource, BaseAnswerResource):
                 ) = processor.resume_from_tool_actions(
                     data["tool_actions"], data["conversation_id"]
                 )
+                if not processor.decoded_token:
+                    return make_response({"error": "Unauthorized"}, 401)
+                if error := self.check_usage(processor.agent_config):
+                    return error
                 stream = self.complete_stream(
                     question="",
                     agent=agent,
diff --git a/application/api/answer/routes/stream.py b/application/api/answer/routes/stream.py
index 2a6b9a11..9bfe5241 100644
--- a/application/api/answer/routes/stream.py
+++ b/application/api/answer/routes/stream.py
@@ -92,6 +92,14 @@ class StreamResource(Resource, BaseAnswerResource):
                 ) = processor.resume_from_tool_actions(
                     data["tool_actions"], data["conversation_id"]
                 )
+                if not processor.decoded_token:
+                    return Response(
+                        self.error_stream_generate("Unauthorized"),
+                        status=401,
+                        mimetype="text/event-stream",
+                    )
+                if error := self.check_usage(processor.agent_config):
+                    return error
                 return Response(
                     self.complete_stream(
                         question="",
diff --git a/tests/agents/test_mcp_tool.py b/tests/agents/test_mcp_tool.py
index eb17c757..c3d4d45b 100644
--- a/tests/agents/test_mcp_tool.py
+++ b/tests/agents/test_mcp_tool.py
@@ -28,6 +28,7 @@ def _patch_mcp_globals(monkeypatch):
     monkeypatch.setattr(mcp_mod, "mongo", mock_mongo)
     monkeypatch.setattr(mcp_mod, "db", mock_db)
     monkeypatch.setattr(mcp_mod, "_mcp_clients_cache", {})
+    monkeypatch.setattr(mcp_mod, "validate_url", lambda url: url)
 
 
 @pytest.fixture