feat: wire AgentLoop to DisplayAdapter and async permissions
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 ---
|
||||
|
||||
Reference in New Issue
Block a user