v2.3.0 (#429)
Co-authored-by: Carlo <carloantonio.patti@mistral.ai> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Clement Sirieix <clem.sirieix@gmail.com> Co-authored-by: Michel Thomazo <michel.thomazo@mistral.ai> Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-authored-by: Thomas Kenbeek <thomas.kenbeek@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin
GitHub
parent
a560a47ce8
commit
5d2e01a6d7
139 changed files with 7152 additions and 1457 deletions
|
|
@ -3,44 +3,67 @@ from __future__ import annotations
|
|||
import pytest
|
||||
|
||||
from tests.conftest import build_test_agent_loop, build_test_vibe_config
|
||||
from vibe.core.agents.models import BUILTIN_AGENTS, AgentProfile, BuiltinAgentName
|
||||
from vibe.core.agents.models import BUILTIN_AGENTS, CHAT, AgentProfile, BuiltinAgentName
|
||||
from vibe.core.config import VibeConfig
|
||||
from vibe.core.middleware import (
|
||||
CHAT_AGENT_EXIT,
|
||||
CHAT_AGENT_REMINDER,
|
||||
PLAN_AGENT_EXIT,
|
||||
PLAN_AGENT_REMINDER,
|
||||
ConversationContext,
|
||||
MiddlewareAction,
|
||||
MiddlewarePipeline,
|
||||
PlanAgentMiddleware,
|
||||
ReadOnlyAgentMiddleware,
|
||||
ResetReason,
|
||||
)
|
||||
from vibe.core.types import AgentStats
|
||||
from vibe.core.types import AgentStats, MessageList
|
||||
|
||||
REMINDER = "test reminder"
|
||||
EXIT_MSG = "test exit"
|
||||
TARGET_AGENT = BuiltinAgentName.PLAN
|
||||
|
||||
|
||||
def _build_middleware(
|
||||
profile_getter,
|
||||
agent_name: str = TARGET_AGENT,
|
||||
reminder: str = REMINDER,
|
||||
exit_message: str = EXIT_MSG,
|
||||
) -> ReadOnlyAgentMiddleware:
|
||||
return ReadOnlyAgentMiddleware(profile_getter, agent_name, reminder, exit_message)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ctx(vibe_config: VibeConfig) -> ConversationContext:
|
||||
return ConversationContext(messages=[], stats=AgentStats(), config=vibe_config)
|
||||
return ConversationContext(
|
||||
messages=MessageList(), stats=AgentStats(), config=vibe_config
|
||||
)
|
||||
|
||||
|
||||
class TestPlanAgentMiddleware:
|
||||
class TestReadOnlyAgentMiddleware:
|
||||
@pytest.mark.asyncio
|
||||
async def test_injects_reminder_when_plan_agent_active(
|
||||
async def test_injects_reminder_when_target_agent_active(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
assert result.message == REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_inject_when_default_agent(
|
||||
self, ctx: ConversationContext
|
||||
@pytest.mark.parametrize(
|
||||
"agent_name",
|
||||
[
|
||||
BuiltinAgentName.DEFAULT,
|
||||
BuiltinAgentName.AUTO_APPROVE,
|
||||
BuiltinAgentName.ACCEPT_EDITS,
|
||||
],
|
||||
)
|
||||
async def test_does_not_inject_when_non_target_agent(
|
||||
self, ctx: ConversationContext, agent_name: str
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
)
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[agent_name])
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
|
|
@ -48,89 +71,56 @@ class TestPlanAgentMiddleware:
|
|||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_inject_when_auto_approve_agent(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_inject_when_accept_edits_agent(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.ACCEPT_EDITS]
|
||||
)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_injects_reminder_only_once_while_in_plan_mode(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
async def test_injects_reminder_only_once(self, ctx: ConversationContext) -> None:
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
|
||||
result1 = await middleware.before_turn(ctx)
|
||||
assert result1.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result1.message == PLAN_AGENT_REMINDER
|
||||
assert result1.message == REMINDER
|
||||
|
||||
result2 = await middleware.before_turn(ctx)
|
||||
assert result2.action == MiddlewareAction.CONTINUE
|
||||
assert result2.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_injects_exit_message_when_leaving_plan_mode(
|
||||
async def test_injects_exit_message_when_leaving(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
# Enter plan mode
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
# Leave plan mode
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
assert result.message == EXIT_MSG
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reinjects_reminder_when_reentering_plan_mode(
|
||||
async def test_reinjects_reminder_on_reentry(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
# Enter plan mode - should inject reminder
|
||||
result1 = await middleware.before_turn(ctx)
|
||||
assert result1.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result1.message == PLAN_AGENT_REMINDER
|
||||
assert result1.message == REMINDER
|
||||
|
||||
# Leave plan mode - should inject exit message
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result2 = await middleware.before_turn(ctx)
|
||||
assert result2.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result2.message == PLAN_AGENT_EXIT
|
||||
assert result2.message == EXIT_MSG
|
||||
|
||||
# Re-enter plan mode - should inject reminder again
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result3 = await middleware.before_turn(ctx)
|
||||
assert result3.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result3.message == PLAN_AGENT_REMINDER
|
||||
assert result3.message == REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_reminder(self, ctx: ConversationContext) -> None:
|
||||
custom_reminder = "Custom plan agent reminder"
|
||||
middleware = PlanAgentMiddleware(
|
||||
custom_reminder = "Custom reminder"
|
||||
middleware = _build_middleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN], reminder=custom_reminder
|
||||
)
|
||||
|
||||
|
|
@ -138,246 +128,216 @@ class TestPlanAgentMiddleware:
|
|||
|
||||
assert result.message == custom_reminder
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_clears_state(self, ctx: ConversationContext) -> None:
|
||||
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
await middleware.before_turn(ctx) # Enter and inject
|
||||
|
||||
middleware.reset()
|
||||
|
||||
# Should inject again after reset
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exit_message_fires_only_once(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
# Enter plan mode
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
# Leave plan mode - first call should inject exit
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
|
||||
# Subsequent calls in default mode should be CONTINUE
|
||||
result2 = await middleware.before_turn(ctx)
|
||||
assert result2.action == MiddlewareAction.CONTINUE
|
||||
assert result2.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_turns_in_plan_mode_after_entry(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
|
||||
# First turn: inject reminder
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
|
||||
# Several more turns in plan mode: all should be CONTINUE
|
||||
for _ in range(5):
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_turns_in_default_mode_after_exit(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx) # exit plan (fires exit message)
|
||||
|
||||
# Several more turns in default mode: all should be CONTINUE
|
||||
for _ in range(5):
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rapid_toggling_plan_default_multiple_cycles(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
for _ in range(3):
|
||||
# Enter plan
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
|
||||
# Leave plan
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exit_to_non_default_agent(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
|
||||
# Switch to auto_approve (not default)
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exit_to_accept_edits_agent(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
|
||||
# Switch to accept_edits
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.ACCEPT_EDITS]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_switching_between_non_plan_agents(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.ACCEPT_EDITS]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_plan_to_plan_entry(self, ctx: ConversationContext) -> None:
|
||||
"""Starting in a non-plan agent then entering plan should inject reminder."""
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
# Now switch to plan
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_while_in_default_after_exiting_plan(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx) # exit plan
|
||||
|
||||
middleware.reset()
|
||||
|
||||
# Still in default mode - should CONTINUE (no phantom exit message)
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_while_in_default_then_reenter_plan(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx) # exit plan
|
||||
|
||||
middleware.reset()
|
||||
|
||||
# Re-enter plan after reset
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_with_compact_reason(self, ctx: ConversationContext) -> None:
|
||||
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
await middleware.before_turn(ctx) # enter and inject
|
||||
|
||||
middleware.reset(ResetReason.COMPACT)
|
||||
|
||||
# Should reinject reminder after compact reset
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_exit_message(self, ctx: ConversationContext) -> None:
|
||||
custom_exit = "Custom exit message"
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(
|
||||
middleware = _build_middleware(
|
||||
lambda: current_profile, exit_message=custom_exit
|
||||
)
|
||||
|
||||
await middleware.before_turn(ctx) # enter plan
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.message == custom_exit
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plan_entry_then_immediate_exit_same_not_possible(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
"""Even if profile changes between two calls, each call sees one transition."""
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = PlanAgentMiddleware(lambda: current_profile)
|
||||
async def test_reset_clears_state(self, ctx: ConversationContext) -> None:
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
middleware.reset()
|
||||
|
||||
# First call: entry
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_REMINDER
|
||||
|
||||
# Second call (still plan): no injection
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
@pytest.mark.asyncio
|
||||
async def test_exit_message_fires_only_once(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
# Third call (switched to default): exit
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == PLAN_AGENT_EXIT
|
||||
assert result.message == EXIT_MSG
|
||||
|
||||
result2 = await middleware.before_turn(ctx)
|
||||
assert result2.action == MiddlewareAction.CONTINUE
|
||||
assert result2.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_turns_after_entry(self, ctx: ConversationContext) -> None:
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
|
||||
for _ in range(5):
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_turns_after_exit(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
for _ in range(5):
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rapid_toggling_multiple_cycles(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
for _ in range(3):
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == REMINDER
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == EXIT_MSG
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exit_to_non_default_agent(self, ctx: ConversationContext) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == EXIT_MSG
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_switching_between_non_target_agents(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.ACCEPT_EDITS]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_target_to_target_entry(self, ctx: ConversationContext) -> None:
|
||||
"""Starting in a non-target agent then entering target should inject reminder."""
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_while_inactive_after_exit(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx)
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
middleware.reset()
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_while_inactive_then_reenter(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
await middleware.before_turn(ctx)
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
middleware.reset()
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_with_compact_reason(self, ctx: ConversationContext) -> None:
|
||||
middleware = _build_middleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
||||
await middleware.before_turn(ctx)
|
||||
|
||||
middleware.reset(ResetReason.COMPACT)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == REMINDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entry_then_continuation_then_exit_then_continuation(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
"""Each call sees one transition at a time."""
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
middleware = _build_middleware(lambda: current_profile)
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == REMINDER
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert result.message == EXIT_MSG
|
||||
|
||||
# Fourth call (still default): no injection
|
||||
result = await middleware.before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
|
||||
class TestMiddlewarePipelineWithPlanAgent:
|
||||
class TestMiddlewarePipelineWithReadOnlyAgent:
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_includes_plan_agent_injection(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
async def test_pipeline_includes_injection(self, ctx: ConversationContext) -> None:
|
||||
pipeline = MiddlewarePipeline()
|
||||
pipeline.add(PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN]))
|
||||
pipeline.add(
|
||||
ReadOnlyAgentMiddleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN],
|
||||
BuiltinAgentName.PLAN,
|
||||
PLAN_AGENT_REMINDER,
|
||||
PLAN_AGENT_EXIT,
|
||||
)
|
||||
)
|
||||
|
||||
result = await pipeline.run_before_turn(ctx)
|
||||
|
||||
|
|
@ -385,20 +345,73 @@ class TestMiddlewarePipelineWithPlanAgent:
|
|||
assert PLAN_AGENT_REMINDER in (result.message or "")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_skips_injection_when_not_plan_agent(
|
||||
async def test_pipeline_skips_injection_when_not_target_agent(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
pipeline = MiddlewarePipeline()
|
||||
pipeline.add(
|
||||
PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.DEFAULT])
|
||||
ReadOnlyAgentMiddleware(
|
||||
lambda: BUILTIN_AGENTS[BuiltinAgentName.DEFAULT],
|
||||
BuiltinAgentName.PLAN,
|
||||
PLAN_AGENT_REMINDER,
|
||||
PLAN_AGENT_EXIT,
|
||||
)
|
||||
)
|
||||
|
||||
result = await pipeline.run_before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_direct_plan_to_chat_transition_delivers_both_messages(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
pipeline = MiddlewarePipeline()
|
||||
pipeline.add(
|
||||
ReadOnlyAgentMiddleware(
|
||||
lambda: current_profile,
|
||||
BuiltinAgentName.PLAN,
|
||||
PLAN_AGENT_REMINDER,
|
||||
PLAN_AGENT_EXIT,
|
||||
)
|
||||
)
|
||||
pipeline.add(
|
||||
ReadOnlyAgentMiddleware(
|
||||
lambda: current_profile,
|
||||
BuiltinAgentName.CHAT,
|
||||
CHAT_AGENT_REMINDER,
|
||||
CHAT_AGENT_EXIT,
|
||||
)
|
||||
)
|
||||
|
||||
class TestPlanAgentMiddlewareIntegration:
|
||||
result = await pipeline.run_before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert PLAN_AGENT_REMINDER in (result.message or "")
|
||||
|
||||
current_profile = CHAT
|
||||
result = await pipeline.run_before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert PLAN_AGENT_EXIT in (result.message or "")
|
||||
assert CHAT_AGENT_REMINDER in (result.message or "")
|
||||
|
||||
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
||||
result = await pipeline.run_before_turn(ctx)
|
||||
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
||||
assert CHAT_AGENT_EXIT in (result.message or "")
|
||||
assert PLAN_AGENT_REMINDER in (result.message or "")
|
||||
|
||||
|
||||
def _find_plan_middleware(agent) -> ReadOnlyAgentMiddleware:
|
||||
return next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, ReadOnlyAgentMiddleware)
|
||||
and mw._agent_name == BuiltinAgentName.PLAN
|
||||
)
|
||||
|
||||
|
||||
class TestReadOnlyAgentMiddlewareIntegration:
|
||||
@pytest.mark.asyncio
|
||||
async def test_switch_agent_preserves_middleware_state_for_exit_message(
|
||||
self,
|
||||
|
|
@ -414,11 +427,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
)
|
||||
agent = build_test_agent_loop(config=config, agent_name=BuiltinAgentName.PLAN)
|
||||
|
||||
plan_middleware = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware = _find_plan_middleware(agent)
|
||||
|
||||
ctx = ConversationContext(
|
||||
messages=agent.messages, stats=agent.stats, config=agent.config
|
||||
|
|
@ -429,11 +438,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
|
||||
await agent.switch_agent(BuiltinAgentName.DEFAULT)
|
||||
|
||||
plan_middleware_after = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware_after = _find_plan_middleware(agent)
|
||||
assert plan_middleware is plan_middleware_after
|
||||
|
||||
ctx = ConversationContext(
|
||||
|
|
@ -456,11 +461,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
)
|
||||
agent = build_test_agent_loop(config=config, agent_name=BuiltinAgentName.PLAN)
|
||||
|
||||
plan_middleware = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware = _find_plan_middleware(agent)
|
||||
|
||||
ctx = ConversationContext(
|
||||
messages=agent.messages, stats=agent.stats, config=agent.config
|
||||
|
|
@ -497,11 +498,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
)
|
||||
agent = build_test_agent_loop(config=config, agent_name=BuiltinAgentName.PLAN)
|
||||
|
||||
plan_middleware = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware = _find_plan_middleware(agent)
|
||||
|
||||
ctx = ConversationContext(
|
||||
messages=agent.messages, stats=agent.stats, config=agent.config
|
||||
|
|
@ -532,11 +529,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
config=config, agent_name=BuiltinAgentName.DEFAULT
|
||||
)
|
||||
|
||||
plan_middleware = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware = _find_plan_middleware(agent)
|
||||
|
||||
ctx = ConversationContext(
|
||||
messages=agent.messages, stats=agent.stats, config=agent.config
|
||||
|
|
@ -566,11 +559,7 @@ class TestPlanAgentMiddlewareIntegration:
|
|||
)
|
||||
agent = build_test_agent_loop(config=config, agent_name=BuiltinAgentName.PLAN)
|
||||
|
||||
plan_middleware = next(
|
||||
mw
|
||||
for mw in agent.middleware_pipeline.middlewares
|
||||
if isinstance(mw, PlanAgentMiddleware)
|
||||
)
|
||||
plan_middleware = _find_plan_middleware(agent)
|
||||
|
||||
def _ctx():
|
||||
return ConversationContext(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue