Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: guard tool execution against non-compliant API gateway injection

Browse Source
This commit is contained in:
Subal
2026-04-17 20:39:46 +08:00
committed by Xubin Ren
parent aabc3d5017
commit 322da6ca06
4 changed files with 218 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/runner.py
+110 -53
View File
@@ -40,14 +40,20 @@ _MAX_INJECTION_CYCLES = 5
_SNIP_SAFETY_BUFFER = 1024
_MICROCOMPACT_KEEP_RECENT = 10
_MICROCOMPACT_MIN_CHARS = 500
_COMPACTABLE_TOOLS = frozenset({
"read_file", "exec", "grep", "glob",
"web_search", "web_fetch", "list_dir",
})
_COMPACTABLE_TOOLS = frozenset(
{
"read_file",
"exec",
"grep",
"glob",
"web_search",
"web_fetch",
"list_dir",
}
)
_BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]"
@dataclass(slots=True)
class AgentRunSpec:
"""Configuration for a single agent execution."""
@@ -120,11 +126,7 @@ class AgentRunner:
) -> None:
"""Append injected user messages while preserving role alternation."""
for injection in injections:
if (
messages
and injection.get("role") == "user"
and messages[-1].get("role") == "user"
):
if messages and injection.get("role") == "user" and messages[-1].get("role") == "user":
merged = dict(messages[-1])
merged["content"] = cls._merge_message_content(
merged.get("content"),
@@ -174,7 +176,10 @@ class AgentRunner:
self._append_injected_messages(messages, injections)
logger.info(
"Injected {} follow-up message(s) {} ({}/{})",
len(injections), phase, injection_cycles, _MAX_INJECTION_CYCLES,
len(injections),
phase,
injection_cycles,
_MAX_INJECTION_CYCLES,
)
return True, injection_cycles
@@ -190,12 +195,9 @@ class AgentRunner:
return []
try:
signature = inspect.signature(spec.injection_callback)
accepts_limit = (
"limit" in signature.parameters
or any(
parameter.kind is inspect.Parameter.VAR_KEYWORD
for parameter in signature.parameters.values()
)
accepts_limit = "limit" in signature.parameters or any(
parameter.kind is inspect.Parameter.VAR_KEYWORD
for parameter in signature.parameters.values()
)
if accepts_limit:
items = await spec.injection_callback(limit=_MAX_INJECTIONS_PER_TURN)
@@ -218,7 +220,9 @@ class AgentRunner:
dropped = len(injected_messages) - _MAX_INJECTIONS_PER_TURN
logger.warning(
"Injection callback returned {} messages, capping to {} ({} dropped)",
len(injected_messages), _MAX_INJECTIONS_PER_TURN, dropped,
len(injected_messages),
_MAX_INJECTIONS_PER_TURN,
dropped,
)
injected_messages = injected_messages[:_MAX_INJECTIONS_PER_TURN]
return injected_messages
@@ -273,7 +277,7 @@ class AgentRunner:
context.tool_calls = list(response.tool_calls)
self._accumulate_usage(usage, raw_usage)
if response.has_tool_calls:
if response.should_execute_tools:
if hook.wants_streaming():
await hook.on_stream_end(context, resuming=True)
@@ -293,7 +297,9 @@ class AgentRunner:
"model": spec.model,
"assistant_message": assistant_message,
"completed_tool_results": [],
"pending_tool_calls": [tc.to_openai_tool_call() for tc in response.tool_calls],
"pending_tool_calls": [
tc.to_openai_tool_call() for tc in response.tool_calls
],
},
)
@@ -332,7 +338,10 @@ class AgentRunner:
context.stop_reason = stop_reason
await hook.after_iteration(context)
should_continue, injection_cycles = await self._try_drain_injections(
spec, messages, None, injection_cycles,
spec,
messages,
None,
injection_cycles,
phase="after tool error",
)
if should_continue:
@@ -354,7 +363,10 @@ class AgentRunner:
length_recovery_count = 0
# Checkpoint 1: drain injections after tools, before next LLM call
_drained, injection_cycles = await self._try_drain_injections(
spec, messages, None, injection_cycles,
spec,
messages,
None,
injection_cycles,
phase="after tool execution",
)
if _drained:
@@ -362,6 +374,13 @@ class AgentRunner:
await hook.after_iteration(context)
continue
elif response.has_tool_calls:
logger.warning(
"Ignoring tool calls under finish_reason='%s' for %s",
response.finish_reason,
spec.session_key or "default",
)
clean = hook.finalize_content(context, response.content)
if response.finish_reason != "error" and is_blank_text(clean):
empty_content_retries += 1
@@ -406,11 +425,13 @@ class AgentRunner:
)
if hook.wants_streaming():
await hook.on_stream_end(context, resuming=True)
messages.append(build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
))
messages.append(
build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
)
messages.append(build_length_recovery_message())
await hook.after_iteration(context)
continue
@@ -427,7 +448,10 @@ class AgentRunner:
# If injections are found we keep the stream alive (resuming=True)
# so streaming channels don't prematurely finalize the card.
should_continue, injection_cycles = await self._try_drain_injections(
spec, messages, assistant_message, injection_cycles,
spec,
messages,
assistant_message,
injection_cycles,
phase="after final response",
iteration=iteration,
)
@@ -451,7 +475,10 @@ class AgentRunner:
context.stop_reason = stop_reason
await hook.after_iteration(context)
should_continue, injection_cycles = await self._try_drain_injections(
spec, messages, None, injection_cycles,
spec,
messages,
None,
injection_cycles,
phase="after LLM error",
)
if should_continue:
@@ -468,7 +495,10 @@ class AgentRunner:
context.stop_reason = stop_reason
await hook.after_iteration(context)
should_continue, injection_cycles = await self._try_drain_injections(
spec, messages, None, injection_cycles,
spec,
messages,
None,
injection_cycles,
phase="after empty response",
)
if should_continue:
@@ -476,11 +506,14 @@ class AgentRunner:
continue
break
messages.append(assistant_message or build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
))
messages.append(
assistant_message
or build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
)
await self._emit_checkpoint(
spec,
{
@@ -516,7 +549,10 @@ class AgentRunner:
# We ignore should_continue here because the for-loop has already
# exhausted all iterations.
drained_after_max_iterations, injection_cycles = await self._try_drain_injections(
spec, messages, None, injection_cycles,
spec,
messages,
None,
injection_cycles,
phase="after max_iterations",
)
if drained_after_max_iterations:
@@ -568,6 +604,7 @@ class AgentRunner:
tools=spec.tools.get_definitions(),
)
if hook.wants_streaming():
async def _stream(delta: str) -> None:
await hook.on_stream(context, delta)
@@ -621,13 +658,19 @@ class AgentRunner:
tool_results: list[tuple[Any, dict[str, str], BaseException | None]] = []
for batch in batches:
if spec.concurrent_tools and len(batch) > 1:
tool_results.extend(await asyncio.gather(*(
self._run_tool(spec, tool_call, external_lookup_counts)
for tool_call in batch
)))
tool_results.extend(
await asyncio.gather(
*(
self._run_tool(spec, tool_call, external_lookup_counts)
for tool_call in batch
)
)
)
else:
for tool_call in batch:
tool_results.append(await self._run_tool(spec, tool_call, external_lookup_counts))
tool_results.append(
await self._run_tool(spec, tool_call, external_lookup_counts)
)
results: list[Any] = []
events: list[dict[str, str]] = []
@@ -675,7 +718,11 @@ class AgentRunner:
"status": "error",
"detail": prep_error.split(": ", 1)[-1][:120],
}
return prep_error + _HINT, event, RuntimeError(prep_error) if spec.fail_on_tool_error else None
return (
prep_error + _HINT,
event,
RuntimeError(prep_error) if spec.fail_on_tool_error else None,
)
try:
if tool is not None:
result = await tool.execute(**params)
@@ -737,7 +784,11 @@ class AgentRunner:
@staticmethod
def _append_model_error_placeholder(messages: list[dict[str, Any]]) -> None:
if messages and messages[-1].get("role") == "assistant" and not messages[-1].get("tool_calls"):
if (
messages
and messages[-1].get("role") == "assistant"
and not messages[-1].get("tool_calls")
):
return
messages.append(build_assistant_message(_PERSISTED_MODEL_ERROR_PLACEHOLDER))
@@ -827,12 +878,15 @@ class AgentRunner:
insert_at = assistant_idx + 1 + offset
while insert_at < len(updated) and updated[insert_at].get("role") == "tool":
insert_at += 1
updated.insert(insert_at, {
"role": "tool",
"tool_call_id": call_id,
"name": name,
"content": _BACKFILL_CONTENT,
})
updated.insert(
insert_at,
{
"role": "tool",
"tool_call_id": call_id,
"name": name,
"content": _BACKFILL_CONTENT,
},
)
offset += 1
return updated
@@ -891,9 +945,13 @@ class AgentRunner:
if not messages or not spec.context_window_tokens:
return messages
provider_max_tokens = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096)
max_output = spec.max_tokens if isinstance(spec.max_tokens, int) else (
provider_max_tokens if isinstance(provider_max_tokens, int) else 4096
provider_max_tokens = getattr(
getattr(self.provider, "generation", None), "max_tokens", 4096
)
max_output = (
spec.max_tokens
if isinstance(spec.max_tokens, int)
else (provider_max_tokens if isinstance(provider_max_tokens, int) else 4096)
)
budget = spec.context_block_limit or (
spec.context_window_tokens - max_output - _SNIP_SAFETY_BUFFER
@@ -976,4 +1034,3 @@ class AgentRunner:
if current:
batches.append(current)
return batches
nanobot/heartbeat/service.py
+22 -8
View File
@@ -93,18 +93,29 @@ class HeartbeatService:
response = await self.provider.chat_with_retry(
messages=[
{"role": "system", "content": "You are a heartbeat agent. Call the heartbeat tool to report your decision."},
{"role": "user", "content": (
f"Current Time: {current_time_str(self.timezone)}\n\n"
"Review the following HEARTBEAT.md and decide whether there are active tasks.\n\n"
f"{content}"
)},
{
"role": "system",
"content": "You are a heartbeat agent. Call the heartbeat tool to report your decision.",
},
{
"role": "user",
"content": (
f"Current Time: {current_time_str(self.timezone)}\n\n"
"Review the following HEARTBEAT.md and decide whether there are active tasks.\n\n"
f"{content}"
),
},
],
tools=_HEARTBEAT_TOOL,
model=self.model,
)
if not response.has_tool_calls:
if not response.should_execute_tools:
if response.has_tool_calls:
logger.warning(
"Ignoring tool calls under finish_reason='%s' in heartbeat",
response.finish_reason,
)
return "skip", ""
args = response.tool_calls[0].arguments
@@ -166,7 +177,10 @@ class HeartbeatService:
if response:
should_notify = await evaluate_response(
response, tasks, self.provider, self.model,
response,
tasks,
self.provider,
self.model,
)
if should_notify and self.on_notify:
logger.info("Heartbeat: completed, delivering response")
nanobot/providers/base.py
+68 -37
View File
@@ -18,6 +18,7 @@ from nanobot.utils.helpers import image_placeholder_text
@dataclass
class ToolCallRequest:
"""A tool call request from the LLM."""
id: str
name: str
arguments: dict[str, Any]
@@ -40,13 +41,16 @@ class ToolCallRequest:
if self.provider_specific_fields:
tool_call["provider_specific_fields"] = self.provider_specific_fields
if self.function_provider_specific_fields:
tool_call["function"]["provider_specific_fields"] = self.function_provider_specific_fields
tool_call["function"]["provider_specific_fields"] = (
self.function_provider_specific_fields
)
return tool_call
@dataclass
class LLMResponse:
"""Response from an LLM provider."""
content: str | None
tool_calls: list[ToolCallRequest] = field(default_factory=list)
finish_reason: str = "stop"
@@ -67,6 +71,18 @@ class LLMResponse:
"""Check if response contains tool calls."""
return len(self.tool_calls) > 0
@property
def should_execute_tools(self) -> bool:
"""Check if tool calls should be executed (guards against gateway injection).
Only execute when finish_reason explicitly signals tool intent.
Tool calls under any other finish_reason (refusal, content_filter, error, etc.)
are treated as anomalous and should not be executed.
"""
if not self.has_tool_calls:
return False
return self.finish_reason == "tool_calls"
@dataclass(frozen=True)
class GenerationSettings:
@@ -103,24 +119,28 @@ class LLMProvider(ABC):
)
_RETRYABLE_STATUS_CODES = frozenset({408, 409, 429})
_TRANSIENT_ERROR_KINDS = frozenset({"timeout", "connection"})
_NON_RETRYABLE_429_ERROR_TOKENS = frozenset({
"insufficient_quota",
"quota_exceeded",
"quota_exhausted",
"billing_hard_limit_reached",
"insufficient_balance",
"credit_balance_too_low",
"billing_not_active",
"payment_required",
})
_RETRYABLE_429_ERROR_TOKENS = frozenset({
"rate_limit_exceeded",
"rate_limit_error",
"too_many_requests",
"request_limit_exceeded",
"requests_limit_exceeded",
"overloaded_error",
})
_NON_RETRYABLE_429_ERROR_TOKENS = frozenset(
{
"insufficient_quota",
"quota_exceeded",
"quota_exhausted",
"billing_hard_limit_reached",
"insufficient_balance",
"credit_balance_too_low",
"billing_not_active",
"payment_required",
}
)
_RETRYABLE_429_ERROR_TOKENS = frozenset(
{
"rate_limit_exceeded",
"rate_limit_error",
"too_many_requests",
"request_limit_exceeded",
"requests_limit_exceeded",
"overloaded_error",
}
)
_NON_RETRYABLE_429_TEXT_MARKERS = (
"insufficient_quota",
"insufficient quota",
@@ -164,7 +184,11 @@ class LLMProvider(ABC):
if isinstance(content, str) and not content:
clean = dict(msg)
clean["content"] = None if (msg.get("role") == "assistant" and msg.get("tool_calls")) else "(empty)"
clean["content"] = (
None
if (msg.get("role") == "assistant" and msg.get("tool_calls"))
else "(empty)"
)
result.append(clean)
continue
@@ -338,10 +362,7 @@ class LLMProvider(ABC):
def _is_retryable_429_response(cls, response: LLMResponse) -> bool:
type_token = cls._normalize_error_token(response.error_type)
code_token = cls._normalize_error_token(response.error_code)
semantic_tokens = {
token for token in (type_token, code_token)
if token is not None
}
semantic_tokens = {token for token in (type_token, code_token) if token is not None}
if any(token in cls._NON_RETRYABLE_429_ERROR_TOKENS for token in semantic_tokens):
return False
@@ -495,9 +516,13 @@ class LLMProvider(ABC):
streaming should override this method.
"""
response = await self.chat(
messages=messages, tools=tools, model=model,
max_tokens=max_tokens, temperature=temperature,
reasoning_effort=reasoning_effort, tool_choice=tool_choice,
messages=messages,
tools=tools,
model=model,
max_tokens=max_tokens,
temperature=temperature,
reasoning_effort=reasoning_effort,
tool_choice=tool_choice,
)
if on_content_delta and response.content:
await on_content_delta(response.content)
@@ -534,9 +559,13 @@ class LLMProvider(ABC):
reasoning_effort = self.generation.reasoning_effort
kw: dict[str, Any] = dict(
messages=messages, tools=tools, model=model,
max_tokens=max_tokens, temperature=temperature,
reasoning_effort=reasoning_effort, tool_choice=tool_choice,
messages=messages,
tools=tools,
model=model,
max_tokens=max_tokens,
temperature=temperature,
reasoning_effort=reasoning_effort,
tool_choice=tool_choice,
on_content_delta=on_content_delta,
)
return await self._run_with_retry(
@@ -576,9 +605,13 @@ class LLMProvider(ABC):
reasoning_effort = self.generation.reasoning_effort
kw: dict[str, Any] = dict(
messages=messages, tools=tools, model=model,
max_tokens=max_tokens, temperature=temperature,
reasoning_effort=reasoning_effort, tool_choice=tool_choice,
messages=messages,
tools=tools,
model=model,
max_tokens=max_tokens,
temperature=temperature,
reasoning_effort=reasoning_effort,
tool_choice=tool_choice,
)
return await self._run_with_retry(
self._safe_chat,
@@ -706,7 +739,7 @@ class LLMProvider(ABC):
if response.finish_reason != "error":
return response
last_response = response
error_key = ((response.content or "").strip().lower() or None)
error_key = (response.content or "").strip().lower() or None
if error_key and error_key == last_error_key:
identical_error_count += 1
else:
@@ -748,9 +781,7 @@ class LLMProvider(ABC):
(response.content or "")[:120].lower(),
)
if on_retry_wait:
await on_retry_wait(
f"Model request failed after {attempt} retries, giving up."
)
await on_retry_wait(f"Model request failed after {attempt} retries, giving up.")
break
base_delay = delays[min(attempt - 1, len(delays) - 1)]
nanobot/utils/evaluator.py
+18 -8
View File
@@ -39,6 +39,7 @@ _EVALUATE_TOOL = [
}
]
async def evaluate_response(
response: str,
task_context: str,
@@ -55,12 +56,15 @@ async def evaluate_response(
llm_response = await provider.chat_with_retry(
messages=[
{"role": "system", "content": render_template("agent/evaluator.md", part="system")},
{"role": "user", "content": render_template(
"agent/evaluator.md",
part="user",
task_context=task_context,
response=response,
)},
{
"role": "user",
"content": render_template(
"agent/evaluator.md",
part="user",
task_context=task_context,
response=response,
),
},
],
tools=_EVALUATE_TOOL,
model=model,
@@ -68,8 +72,14 @@ async def evaluate_response(
temperature=0.0,
)
if not llm_response.has_tool_calls:
logger.warning("evaluate_response: no tool call returned, defaulting to notify")
if not llm_response.should_execute_tools:
if llm_response.has_tool_calls:
logger.warning(
"evaluate_response: ignoring tool calls under finish_reason='%s', defaulting to notify",
llm_response.finish_reason,
)
else:
logger.warning("evaluate_response: no tool call returned, defaulting to notify")
return True
args = llm_response.tool_calls[0].arguments