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.streaming import StreamHandler
from app.tools.registry import ToolRegistry
from app.utils.display import (
print_error,
print_iteration_header,
print_tool_call,
print_tool_result,
print_token_usage,
print_warning,
)
from app.utils.display import DisplayAdapter
from app.utils.logging import get_logger
logger = get_logger(__name__)
@@ -42,6 +35,7 @@ class AgentLoop:
handler: StreamHandler,
registry: ToolRegistry,
permissions: PermissionsService,
display: DisplayAdapter | None = None,
) -> None:
self._config = config
self._ctx = ctx
@@ -49,6 +43,7 @@ class AgentLoop:
self._handler = handler
self._registry = registry
self._permissions = permissions
self._display = display
self._tools_schema = registry.get_openai_tools_schema()
self._system_prompt = self._build_system_prompt()
self._cancelled = False
@@ -92,7 +87,8 @@ class AgentLoop:
reasoning_only_streak = 0
for iteration in range(1, max_iter + 1):
if self._cancelled:
print_warning("Agent loop cancelled.")
if self._display:
self._display.write_warning("Agent loop cancelled.")
break
# 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)
dropped = self._ctx.truncate_history(system_tokens)
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:
print_warning("Token budget exceeded, cannot truncate further. Stopping.")
if self._display:
self._display.write_warning("Token budget exceeded, cannot truncate further. Stopping.")
break
if iteration > 1:
print_iteration_header(iteration, max_iter)
if self._display:
self._display.write_iteration_header(iteration, max_iter)
# Stream LLM response
assistant_msg = await self._llm_step()
@@ -127,11 +126,11 @@ class AgentLoop:
if 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
if total == 0:
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()
@@ -142,17 +141,19 @@ class AgentLoop:
if reasoning_only_streak >= _MAX_REASONING_RETRIES:
# Nudge the model by injecting a user hint
print_warning(
f"Model produced reasoning but no response {reasoning_only_streak} times. "
"Nudging model to respond..."
)
if self._display:
self._display.write_warning(
f"Model produced reasoning but no response {reasoning_only_streak} times. "
"Nudging model to respond..."
)
self._ctx.add_message(
"user",
"Please respond with your answer. Do not just think — provide your actual response.",
)
reasoning_only_streak = 0
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
# Successful response — reset streak
@@ -163,7 +164,7 @@ class AgentLoop:
break
# 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
for result in results:
@@ -179,7 +180,8 @@ class AgentLoop:
if any(r.tool_name == "finish" for r in results):
break
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:
"""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)
return await self._handler.process_stream(chunk_iter)
except KeyboardInterrupt:
print_warning("Response interrupted.")
if self._display:
self._display.write_warning("Response interrupted.")
self._handler.reset()
return None
except (LLMConnectionError, LLMStreamError) as e:
partial = self._handler.get_partial_message()
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
print_error(f"Connection error: {e}")
if self._display:
self._display.write_error(f"Connection error: {e}")
return None
except LLMError as e:
print_error(f"LLM error: {e}")
if self._display:
self._display.write_error(f"LLM error: {e}")
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.
Args:
@@ -226,8 +232,8 @@ class AgentLoop:
tc_id = tc.id
# Display the tool call
if self._config.display.show_tool_calls:
print_tool_call(name, tc.function.arguments)
if self._config.display.show_tool_calls and self._display:
self._display.write_tool_call(name, tc.function.arguments)
# Parse arguments
try:
@@ -240,8 +246,8 @@ class AgentLoop:
error=f"Invalid JSON in arguments: {e}",
)
results.append(result)
if self._config.display.show_tool_calls:
print_tool_result(name, result.error or "", is_error=True)
if self._config.display.show_tool_calls and self._display:
self._display.write_tool_result(name, result.error or "", is_error=True)
continue
# Look up tool
@@ -254,13 +260,13 @@ class AgentLoop:
error=f"Unknown tool '{name}'. Available: {available_names}",
)
results.append(result)
if self._config.display.show_tool_calls:
print_tool_result(name, result.error or "", is_error=True)
if self._config.display.show_tool_calls and self._display:
self._display.write_tool_result(name, result.error or "", is_error=True)
continue
# Check permissions (truncate args for display in prompt)
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(
tool_call_id=tc_id,
tool_name=name,
@@ -268,17 +274,17 @@ class AgentLoop:
error=f"Permission denied for tool '{name}'",
)
results.append(result)
if self._config.display.show_tool_calls:
print_tool_result(name, result.error or "", is_error=True)
if self._config.display.show_tool_calls and self._display:
self._display.write_tool_result(name, result.error or "", is_error=True)
continue
# Execute tool (BaseTool.run never raises)
result = tool.run(tc_id, parsed_args)
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
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

View File

@@ -21,6 +21,7 @@ from app.services.llm import LLMClient
from app.services.permissions import PermissionsService
from app.services.streaming import StreamHandler
from app.tools.registry import ToolRegistry, create_default_registry
from app.utils.display import DisplayAdapter
@pytest.fixture
@@ -76,6 +77,11 @@ def permissions(config: AppConfig) -> PermissionsService:
return PermissionsService(config.permissions)
@pytest.fixture
def display() -> MagicMock:
return MagicMock(spec=DisplayAdapter)
@pytest.fixture
def agent(
config: AppConfig,
@@ -84,8 +90,9 @@ def agent(
handler: MagicMock,
registry: ToolRegistry,
permissions: PermissionsService,
display: MagicMock,
) -> 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:

View File

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