Agent Shield

A security testing tool for AI chatbots: it runs prompt-injection attacks, flags failures with evidence, suggests defenses, then reruns the same tests to measure whether the fix helped.

Comment

Inspiration

AI chatbots and agents are being shipped into banks, hospitals, and customer support faster than security testing can keep up. The #1 risk on the OWASP Top 10 for LLM Applications is prompt injection: a single crafted message can make a chatbot ignore its rules, leak confidential data, or act outside its role.

Most teams test this by hand, if at all. Someone types a few "ignore previous instructions" prompts, says "looks fine," and ships. There is no repeatable process, no evidence a fix worked, and no way to compare before and after.

We asked one question: what if testing an AI chatbot's defenses worked like a unit test suite? Run it, see what breaks, fix it, run it again, and check whether the fix helped.

That became AgentShield.

What it does

AgentShield is a web-based red-team testing tool for LLM chatbots. It closes the full security loop in one place:

Test → Detect → Explain → Fix → Retest → Compare

  1. Test: runs a fixed suite of controlled, non-destructive attacks (direct prompt injection, instruction override, jailbreaks, confidentiality probes, fabricated-action checks) against a sample chatbot that has explicit rules and a fictional protected secret.
  2. Detect: every reply goes through a hybrid evaluator: deterministic code checks first, then an AI model judge for cases code can't decide.
  3. Explain: each result gets one of four outcomes: passed, confirmed fail, suspected fail, or inconclusive, with the prompt, the response, and the reason.
  4. Fix: for failed tests, the model generates advisory defensive suggestions, such as hardened system-prompt rules. The chatbot's instructions are editable in the UI.
  5. Retest: the same suite runs again against the updated chatbot.
  6. Compare: a before/after view shows which attacks the change actually blocked.

We are explicit about limits: a score only describes performance on our test suite. AgentShield is a developer testing tool, not a certified security assessment.

Project status (honest note)

The full pipeline, evaluator, API and dashboard are built and running. Our inference layer is wired to Nebius Token Factory with NVIDIA Nemotron 3 Super (nvidia/nemotron-3-super-120b-a12b).

However, Token Factory access was not available in our region, so none of our team could get an API key and we could not complete live model runs before the deadline. The dashboard in our demo runs in mock mode and is labeled SAMPLE DATA. Those numbers are illustrative, not real Nemotron results, and we do not present them as such.

How we built it

Architecture

┌──────────────────────────────────────────────────────────────┐
│                  FRONTEND (HTML / CSS / JS)                  │
│  Rules view │ Run Tests │ Results │ Suggestions │ Compare    │
└──────────────────────────────┬───────────────────────────────┘
                               │  REST / JSON (api_contract.md)
┌──────────────────────────────▼───────────────────────────────┐
│                     BACKEND (FastAPI)                        │
│                                                              │
│  ┌────────────┐   ┌──────────────┐   ┌────────────────────┐  │
│  │ test_cases │──▶│  runner.py   │──▶│ target_chatbot.py  │  │
│  │   .json    │   │ orchestrates │   │ rules + fake secret│  │
│  └────────────┘   └──────┬───────┘   └─────────┬──────────┘  │
│                          │ reply                │ inference  │
│                   ┌──────▼───────┐       ┌──────▼──────────┐  │
│                   │ evaluator.py │◀─────▶│ nebius_client   │  │
│                   │ 1. code check│       │ Nebius Token    │  │
│                   │ 2. model     │       │ Factory         │  │
│                   │    judge     │       │ (Nemotron)      │  │
│                   └──────┬───────┘       └──────▲──────────┘  │
│                          │ verdict              │             │
│            ┌─────────────▼───────┐    ┌─────────┴─────────┐   │
│            │ suggestions.py      │    │   scoring.py      │   │
│            │ defensive fixes     │    │ pass rate, delta  │   │
│            └─────────────┬───────┘    └─────────┬─────────┘   │
│                          │                      │             │
│                    ┌─────▼──────────────────────▼─────┐       │
│                    │  db.py  (SQLite: run history)    │       │
│                    └──────────────────────────────────┘       │
└──────────────────────────────────────────────────────────────┘

The closed-loop pipeline

 Attack suite ──▶ Target chatbot ──▶ Hybrid evaluator ──▶ Verdict
      ▲                                                     │
      │                                                     ▼
 Retest (same suite)  ◀── Apply fix ◀── Suggested defenses ◀┘
      │
      ▼
 Before / After comparison

Hybrid evaluation (our core design decision)

Judging LLM output with another LLM alone is noisy and hard to trust. Code alone misses subtle failures. We use both, in order:

Layer Method Handles Why
1. Deterministic Plain code rules, e.g. "does the reply contain the secret string?" Clear-cut leaks Reproducible, no hallucination, fast
2. Model judge LLM returns a strict JSON verdict Ambiguous behavior (partial leaks, rule violations) Catches what code can't
3. Safety net Invalid or unparseable judge output becomes inconclusive Judge failures Never counts a broken judgment as a pass

Code-proven leaks are labeled confirmed fail. Failures found only by the model judge are labeled suspected fail, because the judge can be wrong. This separates what we know from what we suspect.

Tech stack

Layer Technology Purpose
Inference Nebius Token Factory + NVIDIA Nemotron 3 Super (nvidia/nemotron-3-super-120b-a12b) Target chatbot, judge, and suggestions (integration built; live runs blocked by regional access)
Backend Python, FastAPI, Uvicorn, Pydantic API and orchestration
Model client OpenAI-compatible Python client One gateway to Token Factory
Storage SQLite Persist runs for before/after comparison
Frontend HTML, CSS, vanilla JS Dashboard, no build step
Config python-dotenv Secrets stay out of code and Git
Demo mode NEBIUS_MOCK=1 Runs the full UI without a model, shown as SAMPLE DATA

Engineering decisions

  • Single model gateway: every model call goes through nebius_client.py, so swapping models or providers is a configuration change.
  • Editable target prompt: the chatbot's system prompt is a live, editable object, which is what makes "apply fix → retest" possible.
  • API contract first: we agreed the JSON schema (docs/api_contract.md) on Day 1, so backend and frontend were built in parallel.
  • Safe by design: all attacks run only against our own sample chatbot with fictional secrets. Nothing is destructive and nothing touches real systems.
  • Mock mode with clear labeling: the UI shows a SAMPLE DATA badge whenever results are not from a live model.
  • Built by a team of 2 in 7 days, with clear file ownership between AI/Backend and Frontend/Product.

Challenges we ran into

  • No model access: Nebius Token Factory was not available in our region, so no one on the team could get an API key. We built the pipeline against the documented API and added a labeled mock mode so the whole flow could still be developed and demonstrated.
  • Judging an LLM with an LLM: judges can be inconsistent or return malformed output. We used deterministic checks first, strict JSON for the judge, and an inconclusive fallback instead of guessing.
  • Avoiding false confidence: a high pass rate is meaningless if the tests are weak, so the limits are written into the product itself.
  • Fixes must be provable: AI-suggested defenses can sound good and still fail. That is why the retest step exists.
  • Tight timeline: 7 days and 2 people. We cut scope to one complete loop instead of many half-built features.

Accomplishments that we're proud of

  • A complete test → detect → fix → retest loop with a working API and dashboard, not just a scanner.
  • A trustworthy evaluator with four honest outcome states instead of a naive pass/fail.
  • Honest reporting: we label sample data as sample data and separate proven failures from suspected ones.
  • A model-agnostic gateway configured for Nemotron 3 Super on Token Factory, ready for live runs the moment access is available.
  • Responsible design: fictional data only and non-destructive tests.
  • A clean, documented repo with a plan, API contract, MIT license, and reproducible setup.

What we learned

  • Security needs evidence, not vibes. A repeatable suite beats "we tried a few prompts."
  • Deterministic checks beat model judgment wherever they apply. Use the LLM only for what code can't decide.
  • Prompt hardening helps but is not a complete defense. That's why suggestions are advisory and every fix is retested.
  • Contracts unlock speed. Agreeing the API on Day 1 saved us days.
  • Build for access problems. A clean model gateway and a labeled mock mode kept the project moving when the key never arrived.

What's next for Agent Shield

  • Run live on Nebius Token Factory with Nemotron and publish real before/after results.
  • Bring-your-own-agent: point AgentShield at any chatbot or agent endpoint instead of only our sample.
  • Bigger attack library mapped to the OWASP LLM Top 10 categories.
  • Agent and tool-use testing: indirect injection through documents, web pages, and tool outputs.
  • CI/CD integration: a GitHub Action that fails the build when the pass rate drops.
  • Multi-model benchmarking: compare Nemotron variants as target and as judge under identical attacks.

Built With

Share this project:

Updates

Submission history