Co-authored-by: Clément Drouin <clement.drouin@mistral.ai>
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com>
Co-authored-by: Mert Unsal <mertunsal1905@gmail.com>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Val <102326092+vdeva@users.noreply.github.com>
Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com>
Co-authored-by: maximevoisin-pm <maxime.voisin@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Mathias Gesbert 2026-05-27 17:10:40 +02:00 committed by GitHub
parent adb1ca74ce
commit cf3f4ca58f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
143 changed files with 3457 additions and 1351 deletions

View file

@ -14,6 +14,7 @@ from vibe.core.middleware import (
MiddlewarePipeline,
ReadOnlyAgentMiddleware,
ResetReason,
TokenLimitMiddleware,
make_plan_agent_reminder,
)
from vibe.core.types import AgentStats, MessageList
@ -634,3 +635,31 @@ class TestReadOnlyAgentMiddlewareIntegration:
# 8. Stay in default: no injection
r = await plan_middleware.before_turn(_ctx())
assert r.action == MiddlewareAction.CONTINUE
class TestTokenLimitMiddleware:
@pytest.mark.asyncio
async def test_stops_when_session_total_tokens_exceeds_limit(
self, ctx: ConversationContext
) -> None:
middleware = TokenLimitMiddleware(14)
ctx.stats.session_prompt_tokens = 10
ctx.stats.session_completion_tokens = 5
result = await middleware.before_turn(ctx)
assert result.action == MiddlewareAction.STOP
assert result.reason == "Token limit exceeded: 15 > 14"
@pytest.mark.asyncio
async def test_allows_when_session_total_tokens_matches_limit(
self, ctx: ConversationContext
) -> None:
middleware = TokenLimitMiddleware(15)
ctx.stats.session_prompt_tokens = 10
ctx.stats.session_completion_tokens = 5
result = await middleware.before_turn(ctx)
assert result.action == MiddlewareAction.CONTINUE
assert result.reason is None