Skip to content

fix(agents): validate, bound, and sanitize tool output in agent loop (closes #199) - #611

Open
rohanbuilds-ai wants to merge 5 commits into
GenAI-Security-Project:mainfrom
rohanbuilds-ai:fix/agent-tool-output-validation
Open

rohanbuilds-ai wants to merge 5 commits into
GenAI-Security-Project:mainfrom
rohanbuilds-ai:fix/agent-tool-output-validation

Conversation

@rohanbuilds-ai

Copy link
Copy Markdown

Closes #199

Threat Model & Problem Summary

In finbot/agents/base.py (_run_agent_loop), tool execution results were previously serialized via naive json.dumps() / str() fallbacks and appended directly to LLM messages without upper bounds, structure preservation, or content validation.

This introduced three critical risks:

  1. Context Window Exhaustion & Denial of Service: Tools returning oversized outputs (e.g., massive file listings, large API responses, or circular data structures) could exhaust model token budgets or cause client schema crashes.
  2. Chat-Template Delimiter Breakout: Raw control tokens (such as <|im_start|>, <|im_end|>, <|system|>) embedded in external vendor descriptions or documents could escape tool output boundaries in tokenizer expansions.
  3. Indirect Prompt Injection: Data fields controlled by attackers (e.g., malicious vendor names, descriptions, invoice notes) could inject overt instruction overrides (Ignore all previous instructions..., You are now DAN mode...) directly into the LLM context as trusted tool messages.

Solution & Changes

1. Robust Type Normalization & Fallbacks

  • Normalizes None to "{}" and serializes primitives, dicts, and lists safely.
  • Gracefully handles non-serializable objects, circular references, and recursion errors with structured error descriptors rather than crashing the agent loop.

2. Strict Output Length Bounding & Structured Envelope Preservation

  • Introduces AGENT_MAX_TOOL_OUTPUT_LENGTH (default: 25,000 characters; configurable per agent via max_tool_output_length).
  • Structured Outputs (Dict/List): Oversized outputs are wrapped in a bounded, syntactically valid JSON envelope:
    {
      "_truncated": true,
      "_limit": 25000,
      "_notice": "Output exceeded limit of 25000 characters and was truncated.",
      "excerpt": "..."
    }
    Guarantees that downstream parsers (json.loads()) and LLM JSON extractors receive valid JSON rather than broken/unclosed string fragments.
  • Raw Text: Oversized plain text is cleanly truncated with an explicit notice ([Output truncated: X characters exceeded limit of Y]) strictly fitting within max_length.

3. Control-Token Delimiter Neutralization

  • Neutralizes raw chat-template escape delimiters (<|im_start|>, <|im_end|>, <|start_header_id|>, <|end_header_id|>, <|system|>) into safe representations (< |im_start|>) to prevent prompt boundary escaping.

4. Prompt Injection Mitigation & Redacted Telemetry

  • Introduces AGENT_BLOCK_PROMPT_INJECTION: bool = True (configurable globally and per-agent via agent_config["block_prompt_injection"]).
  • When suspicious instruction overrides (ignore previous instructions, DAN mode, system prompt override) are detected:
    • Secure Mode (Default True): Replaces suspicious tool output before it enters LLM messages with [Content blocked: potential prompt injection detected in tool output] (or valid JSON {"error": "...", "blocked": True}).
    • Permissive / Sandbox Mode (False): Permits raw output for educational research or unconstrained sandboxes while still emitting structured telemetry warnings.
  • Telemetry Safety: Emits security logs with category markers (instruction_override, persona_hijack, privilege_escalation) and agent context while strictly redacting raw payloads, vendor records, SSNs, and credentials.

5. CTF Challenge Compatibility

  • FinBot CTF challenges relying on indirect prompt injection (Fine Print, Zero-Click Harvest, Sleeper Agent) teach contextual alignment (e.g. accounting contract amendments, mandatory regulatory filings) rather than overt brute-force markers (ignore previous instructions). Normal business and accounting context is fully preserved.
  • Any challenge or agent requiring raw instruction override passthrough can explicitly set "block_prompt_injection": False on its config without creating a global security bypass for other agents.

Verification & Test Results

  • Tool Output Validation Suite (tests/unit/agents/test_tool_output_validation.py):
    • 29 passed in 2.31s (100% pass rate).
    • Covers type serialization, circular references, length bounding, invalid config fallbacks, delimiter escaping, redacted telemetry, valid JSON truncation envelopes, default prompt-injection blocking in LLM requests, and per-agent permissive overrides.
  • Base Agent Suite (tests/unit/agents/test_base_agent.py):
    • 12 passed in 1.10s (100% pass rate).
  • CTF Backend Suite (tests/unit/ctf/):
    • 27 passed, 1 skipped in 2.27s (100% pass rate).
  • Full Agent Suite (tests/unit/agents/):
    • 80 passed, 15 skipped, 3 failed.
    • The 3 failures are pre-existing baseline failures on upstream main (1450fc4):
      • test_sai_edge_001 & test_sai_edge_002: coroutine type assertions due to upstream _get_user_prompt async signature update without test await.
      • test_sai_von_005: vendor VEND-001 unseeded in test fixture.
  • Linters:
    • black --check: Clean (3 files unchanged).
    • isort --check: Clean (exit code 0).
    • git diff --check: Clean (no whitespace issues).

Known Limitations

  • Heuristic regex scanning targets overt, known instruction-override markers (ignore previous instructions, DAN mode, system prompt override). It is designed as a first-line defense-in-depth filter; complex semantic prompt injections without explicit override markers are mitigated by platform guardrails, tool permissions, and LLM alignment.

…nt overrides

Resolves Issue GenAI-Security-Project#199 acceptance criteria by blocking suspicious instruction-override tool outputs before they enter the LLM context.

- Sets settings.AGENT_BLOCK_PROMPT_INJECTION=True by default.
- Preserves per-agent override via agent_config['block_prompt_injection'].
- Documents regex heuristic boundary vs semantic indirect injections.
- Adds agent loop integration tests verifying payloads do not reach LLM messages.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Issue: Tool output appended to LLM context without length or content validation

1 participant