feat: wire AgentLoop to DisplayAdapter and async permissions

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-11 12:46:28 -05:00
parent cab3fbc1cf
commit 641672e4c7
3 changed files with 76 additions and 52 deletions

View File

@@ -11,14 +11,7 @@ from app.services.llm import LLMClient, LLMConnectionError, LLMError, LLMStreamE
from app.services.permissions import PermissionsService from app.services.permissions import PermissionsService
from app.services.streaming import StreamHandler from app.services.streaming import StreamHandler
from app.tools.registry import ToolRegistry from app.tools.registry import ToolRegistry
from app.utils.display import ( from app.utils.display import DisplayAdapter
print_error,
print_iteration_header,
print_tool_call,
print_tool_result,
print_token_usage,
print_warning,
)
from app.utils.logging import get_logger from app.utils.logging import get_logger
logger = get_logger(__name__) logger = get_logger(__name__)
@@ -42,6 +35,7 @@ class AgentLoop:
handler: StreamHandler, handler: StreamHandler,
registry: ToolRegistry, registry: ToolRegistry,
permissions: PermissionsService, permissions: PermissionsService,
display: DisplayAdapter | None = None,
) -> None: ) -> None:
self._config = config self._config = config
self._ctx = ctx self._ctx = ctx
@@ -49,6 +43,7 @@ class AgentLoop:
self._handler = handler self._handler = handler
self._registry = registry self._registry = registry
self._permissions = permissions self._permissions = permissions
self._display = display
self._tools_schema = registry.get_openai_tools_schema() self._tools_schema = registry.get_openai_tools_schema()
self._system_prompt = self._build_system_prompt() self._system_prompt = self._build_system_prompt()
self._cancelled = False self._cancelled = False
@@ -92,7 +87,8 @@ class AgentLoop:
reasoning_only_streak = 0 reasoning_only_streak = 0
for iteration in range(1, max_iter + 1): for iteration in range(1, max_iter + 1):
if self._cancelled: if self._cancelled:
print_warning("Agent loop cancelled.") if self._display:
self._display.write_warning("Agent loop cancelled.")
break break
# Check token budget — try truncation before giving up # Check token budget — try truncation before giving up
@@ -100,13 +96,16 @@ class AgentLoop:
system_tokens = self._ctx.token_counter.estimate_tokens(self._system_prompt) system_tokens = self._ctx.token_counter.estimate_tokens(self._system_prompt)
dropped = self._ctx.truncate_history(system_tokens) dropped = self._ctx.truncate_history(system_tokens)
if dropped > 0: if dropped > 0:
print_warning(f"Token budget pressure: dropped {dropped} oldest messages.") if self._display:
self._display.write_warning(f"Token budget pressure: dropped {dropped} oldest messages.")
else: else:
print_warning("Token budget exceeded, cannot truncate further. Stopping.") if self._display:
self._display.write_warning("Token budget exceeded, cannot truncate further. Stopping.")
break break
if iteration > 1: if iteration > 1:
print_iteration_header(iteration, max_iter) if self._display:
self._display.write_iteration_header(iteration, max_iter)
# Stream LLM response # Stream LLM response
assistant_msg = await self._llm_step() assistant_msg = await self._llm_step()
@@ -127,11 +126,11 @@ class AgentLoop:
if self._handler.usage: if self._handler.usage:
self._ctx.token_counter.count_usage(self._handler.usage) self._ctx.token_counter.count_usage(self._handler.usage)
if self._config.display.show_token_usage: if self._config.display.show_token_usage and self._display:
total = self._ctx.token_counter.cumulative_usage.total_tokens total = self._ctx.token_counter.cumulative_usage.total_tokens
if total == 0: if total == 0:
total = self._ctx.estimated_tokens total = self._ctx.estimated_tokens
print_token_usage(total, self._ctx.token_counter.budget) self._display.write_token_usage(total, self._ctx.token_counter.budget)
self._handler.reset() self._handler.reset()
@@ -142,17 +141,19 @@ class AgentLoop:
if reasoning_only_streak >= _MAX_REASONING_RETRIES: if reasoning_only_streak >= _MAX_REASONING_RETRIES:
# Nudge the model by injecting a user hint # Nudge the model by injecting a user hint
print_warning( if self._display:
f"Model produced reasoning but no response {reasoning_only_streak} times. " self._display.write_warning(
"Nudging model to respond..." f"Model produced reasoning but no response {reasoning_only_streak} times. "
) "Nudging model to respond..."
)
self._ctx.add_message( self._ctx.add_message(
"user", "user",
"Please respond with your answer. Do not just think — provide your actual response.", "Please respond with your answer. Do not just think — provide your actual response.",
) )
reasoning_only_streak = 0 reasoning_only_streak = 0
else: else:
print_warning("Model produced reasoning but no response. Retrying...") if self._display:
self._display.write_warning("Model produced reasoning but no response. Retrying...")
continue continue
# Successful response — reset streak # Successful response — reset streak
@@ -163,7 +164,7 @@ class AgentLoop:
break break
# Execute tool calls # Execute tool calls
results = self._execute_tool_calls(assistant_msg.tool_calls) results = await self._execute_tool_calls(assistant_msg.tool_calls)
# Add tool results to context # Add tool results to context
for result in results: for result in results:
@@ -179,7 +180,8 @@ class AgentLoop:
if any(r.tool_name == "finish" for r in results): if any(r.tool_name == "finish" for r in results):
break break
else: else:
print_warning(f"Agent reached maximum iterations ({max_iter}). Stopping.") if self._display:
self._display.write_warning(f"Agent reached maximum iterations ({max_iter}). Stopping.")
async def _llm_step(self) -> Message | None: async def _llm_step(self) -> Message | None:
"""Stream one LLM response and return the accumulated Message. """Stream one LLM response and return the accumulated Message.
@@ -195,21 +197,25 @@ class AgentLoop:
chunk_iter = self._client.stream_chat_with_retry(messages, tools=self._tools_schema) chunk_iter = self._client.stream_chat_with_retry(messages, tools=self._tools_schema)
return await self._handler.process_stream(chunk_iter) return await self._handler.process_stream(chunk_iter)
except KeyboardInterrupt: except KeyboardInterrupt:
print_warning("Response interrupted.") if self._display:
self._display.write_warning("Response interrupted.")
self._handler.reset() self._handler.reset()
return None return None
except (LLMConnectionError, LLMStreamError) as e: except (LLMConnectionError, LLMStreamError) as e:
partial = self._handler.get_partial_message() partial = self._handler.get_partial_message()
if partial is not None: if partial is not None:
print_warning(f"Stream interrupted ({e}), returning partial response.") if self._display:
self._display.write_warning(f"Stream interrupted ({e}), returning partial response.")
return partial return partial
print_error(f"Connection error: {e}") if self._display:
self._display.write_error(f"Connection error: {e}")
return None return None
except LLMError as e: except LLMError as e:
print_error(f"LLM error: {e}") if self._display:
self._display.write_error(f"LLM error: {e}")
return None return None
def _execute_tool_calls(self, tool_calls: list[ToolCall]) -> list[ToolResult]: async def _execute_tool_calls(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
"""Execute a list of tool calls with permission checks. """Execute a list of tool calls with permission checks.
Args: Args:
@@ -226,8 +232,8 @@ class AgentLoop:
tc_id = tc.id tc_id = tc.id
# Display the tool call # Display the tool call
if self._config.display.show_tool_calls: if self._config.display.show_tool_calls and self._display:
print_tool_call(name, tc.function.arguments) self._display.write_tool_call(name, tc.function.arguments)
# Parse arguments # Parse arguments
try: try:
@@ -240,8 +246,8 @@ class AgentLoop:
error=f"Invalid JSON in arguments: {e}", error=f"Invalid JSON in arguments: {e}",
) )
results.append(result) results.append(result)
if self._config.display.show_tool_calls: if self._config.display.show_tool_calls and self._display:
print_tool_result(name, result.error or "", is_error=True) self._display.write_tool_result(name, result.error or "", is_error=True)
continue continue
# Look up tool # Look up tool
@@ -254,13 +260,13 @@ class AgentLoop:
error=f"Unknown tool '{name}'. Available: {available_names}", error=f"Unknown tool '{name}'. Available: {available_names}",
) )
results.append(result) results.append(result)
if self._config.display.show_tool_calls: if self._config.display.show_tool_calls and self._display:
print_tool_result(name, result.error or "", is_error=True) self._display.write_tool_result(name, result.error or "", is_error=True)
continue continue
# Check permissions (truncate args for display in prompt) # Check permissions (truncate args for display in prompt)
desc = tc.function.arguments[:120] + "..." if len(tc.function.arguments) > 120 else tc.function.arguments desc = tc.function.arguments[:120] + "..." if len(tc.function.arguments) > 120 else tc.function.arguments
if not self._permissions.check(name, description=desc): if not await self._permissions.check(name, description=desc):
result = ToolResult( result = ToolResult(
tool_call_id=tc_id, tool_call_id=tc_id,
tool_name=name, tool_name=name,
@@ -268,17 +274,17 @@ class AgentLoop:
error=f"Permission denied for tool '{name}'", error=f"Permission denied for tool '{name}'",
) )
results.append(result) results.append(result)
if self._config.display.show_tool_calls: if self._config.display.show_tool_calls and self._display:
print_tool_result(name, result.error or "", is_error=True) self._display.write_tool_result(name, result.error or "", is_error=True)
continue continue
# Execute tool (BaseTool.run never raises) # Execute tool (BaseTool.run never raises)
result = tool.run(tc_id, parsed_args) result = tool.run(tc_id, parsed_args)
results.append(result) results.append(result)
if self._config.display.show_tool_calls: if self._config.display.show_tool_calls and self._display:
is_error = result.status == ToolResultStatus.ERROR is_error = result.status == ToolResultStatus.ERROR
output = result.error if is_error else result.output output = result.error if is_error else result.output
print_tool_result(name, output or "", is_error=is_error) self._display.write_tool_result(name, output or "", is_error=is_error)
return results return results

View File

@@ -21,6 +21,7 @@ from app.services.llm import LLMClient
from app.services.permissions import PermissionsService from app.services.permissions import PermissionsService
from app.services.streaming import StreamHandler from app.services.streaming import StreamHandler
from app.tools.registry import ToolRegistry, create_default_registry from app.tools.registry import ToolRegistry, create_default_registry
from app.utils.display import DisplayAdapter
@pytest.fixture @pytest.fixture
@@ -76,6 +77,11 @@ def permissions(config: AppConfig) -> PermissionsService:
return PermissionsService(config.permissions) return PermissionsService(config.permissions)
@pytest.fixture
def display() -> MagicMock:
return MagicMock(spec=DisplayAdapter)
@pytest.fixture @pytest.fixture
def agent( def agent(
config: AppConfig, config: AppConfig,
@@ -84,8 +90,9 @@ def agent(
handler: MagicMock, handler: MagicMock,
registry: ToolRegistry, registry: ToolRegistry,
permissions: PermissionsService, permissions: PermissionsService,
display: MagicMock,
) -> AgentLoop: ) -> AgentLoop:
return AgentLoop(config, ctx, client, handler, registry, permissions) return AgentLoop(config, ctx, client, handler, registry, permissions, display)
def _make_text_message(content: str) -> Message: def _make_text_message(content: str) -> Message:

View File

@@ -1,7 +1,6 @@
"""Tests for the tool framework and core tools (Phase 4).""" """Tests for the tool framework and core tools (Phase 4)."""
from pathlib import Path from pathlib import Path
from unittest.mock import patch
import pytest import pytest
@@ -48,28 +47,40 @@ class TestBaseTool:
class TestPermissionsService: class TestPermissionsService:
def test_deny_list_blocks(self) -> None: @pytest.mark.asyncio
async def test_deny_list_blocks(self) -> None:
svc = PermissionsService(PermissionsConfig(deny=["dangerous_tool"])) svc = PermissionsService(PermissionsConfig(deny=["dangerous_tool"]))
assert svc.check("dangerous_tool") is False assert await svc.check("dangerous_tool") is False
def test_auto_approve_allows(self) -> None: @pytest.mark.asyncio
async def test_auto_approve_allows(self) -> None:
svc = PermissionsService(PermissionsConfig(auto_approve=["read_file"])) svc = PermissionsService(PermissionsConfig(auto_approve=["read_file"]))
assert svc.check("read_file") is True assert await svc.check("read_file") is True
@patch("app.services.permissions.Confirm.ask", return_value=True) @pytest.mark.asyncio
def test_prompt_user_approved(self, mock_ask: object) -> None: async def test_prompt_callback_approved(self) -> None:
svc = PermissionsService(PermissionsConfig(prompt_user=["write_file"])) svc = PermissionsService(PermissionsConfig(prompt_user=["write_file"]))
assert svc.check("write_file") is True
@patch("app.services.permissions.Confirm.ask", return_value=False) async def approve(tool_name: str, description: str) -> bool:
def test_prompt_user_denied(self, mock_ask: object) -> None: return True
svc.set_prompt_callback(approve)
assert await svc.check("write_file") is True
@pytest.mark.asyncio
async def test_prompt_callback_denied(self) -> None:
svc = PermissionsService(PermissionsConfig(prompt_user=["write_file"])) svc = PermissionsService(PermissionsConfig(prompt_user=["write_file"]))
assert svc.check("write_file") is False
@patch("app.services.permissions.Confirm.ask", return_value=False) async def deny(tool_name: str, description: str) -> bool:
def test_unlisted_tool_prompts(self, mock_ask: object) -> None: return False
svc.set_prompt_callback(deny)
assert await svc.check("write_file") is False
@pytest.mark.asyncio
async def test_unlisted_tool_no_callback_denied(self) -> None:
svc = PermissionsService(PermissionsConfig()) svc = PermissionsService(PermissionsConfig())
assert svc.check("unknown_tool") is False assert await svc.check("unknown_tool") is False
# --- ToolRegistry --- # --- ToolRegistry ---