Clean cruft files, fix README links after docs/ reorg, integrate Cursor's inference enhancements

- Delete notebook_builder.py (empty placeholder) and generate_submission_plots.py (superseded by generate_plots.py)
- Update README links pointing to moved blog files in docs/blog/
- Cursor edits to inference.py / models.py: added Pydantic-validated action parsing, cumulative reward + grader-breakdown logging, optional INFERENCE_SUMMARY_FILE, interactive mode flag — core env contract preserved
- Add conftest.py for pytest discovery

.env.example

@@ -1,7 +1,25 @@
 # Required for inference.py (baseline agent)
-API_BASE_URL=https://router.huggingface.co/v1
-MODEL_NAME=meta-llama/Llama-3.3-70B-Instruct
-HF_TOKEN=your-huggingface-token-here
+# HF_TOKEN is passed as api_key to the OpenAI client — use the token for
+# whatever provider matches API_BASE_URL (Hugging Face, OpenRouter, etc.).
 
-# Optional: point inference at a different env URL
+# Option A — Hugging Face router (defaults in code if env unset)
+# API_BASE_URL=https://router.huggingface.co/v1
+# MODEL_NAME=meta-llama/Llama-3.3-70B-Instruct
+# HF_TOKEN=your-huggingface-token-here
+
+# Option B — OpenRouter (e.g. Meta Llama 3.2 3B Instruct free)
+# API_BASE_URL=https://openrouter.ai/api/v1
+# MODEL_NAME=meta-llama/llama-3.2-3b-instruct:free
+# HF_TOKEN=your-openrouter-key-here
+
+# Optional: point inference at the local env
 ENV_URL=http://localhost:8000
+
+# Optional: append every raw LLM response to a file (debug / audit)
+# INFERENCE_LOG_LLM=outputs/llm_raw.log
+# Optional: append API errors (e.g. 429) when the client falls back to list_tools
+# INFERENCE_LOG_API=outputs/api_errors.log
+# Optional: wait for Enter between steps to avoid rate limits (or use: python inference.py -i)
+# INFERENCE_INTERACTIVE=1
+# INFERENCE_PAUSE=step   # step = after each env step; scenario = only between easy/medium/hard
+# INFERENCE_SUMMARY_FILE=outputs/reward_grader_summary.txt

README.md

@@ -422,9 +422,9 @@ Key research validating our design:
 ## Links & Resources
 
 - **Live Environment**: https://huggingface.co/spaces/Sayuj63/Vapt-env
-- **Blog Post**: [VAPT env: Teaching AI to Reason About Security](./VAPT_env_BLOG_POST_FINAL.md) — Read the full story
+- **Blog Post**: [VAPT env: Teaching AI to Reason About Security](./docs/blog/VAPT_env_BLOG_POST_FINAL.md) — Read the full story
 - **Training Notebook**: [VAPT_env_RL_Training_Colab.ipynb](./AISHA_RL_Training_Colab.ipynb)
 - **Agent Comparison Script**: [generate_plots.py](./generate_plots.py)
-- **Publication Guide**: [BLOG_PUBLICATION_GUIDE.md](./BLOG_PUBLICATION_GUIDE.md) — How to publish the blog post
+- **Publication Guide**: [BLOG_PUBLICATION_GUIDE.md](./docs/blog/BLOG_PUBLICATION_GUIDE.md) — How to publish the blog post
 - **Team**: Your Team Name — [Your Team Members]
 - **Hackathon**: Meta PyTorch OpenEnv Hackathon India 2026

__init__.py

@@ -5,12 +5,37 @@
 
 """Security Audit Environment — AI-powered VAPT training."""
 
-from .client import SecurityAuditEnv
-from .models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState
+# Pytest (and other tools) may import this file as a top-level module when the repo
+# root is treated as a "package" path. Relative imports need a parent package; fall
+# back to same-directory imports so `pytest` still works. Normal `import security_audit_env`
+# from an install uses the relative branch.
+try:
+    from .client import SecurityAuditEnv
+    from .models import (
+        LLMJsonAction,
+        SecurityAuditAction,
+        SecurityAuditObservation,
+        SecurityAuditState,
+        extract_json_object_from_text,
+        parse_llm_action_text,
+    )
+except ImportError:  # pragma: no cover
+    from client import SecurityAuditEnv
+    from models import (
+        LLMJsonAction,
+        SecurityAuditAction,
+        SecurityAuditObservation,
+        SecurityAuditState,
+        extract_json_object_from_text,
+        parse_llm_action_text,
+    )
 
 __all__ = [
+    "LLMJsonAction",
     "SecurityAuditAction",
     "SecurityAuditObservation",
     "SecurityAuditState",
     "SecurityAuditEnv",
+    "extract_json_object_from_text",
+    "parse_llm_action_text",
 ]

client.py

@@ -10,7 +10,10 @@ from typing import Any, Dict
 from openenv.core import EnvClient
 from openenv.core.client_types import StepResult
 
-from .models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState
+try:
+    from .models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState
+except ImportError:  # pragma: no cover
+    from models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState
 
 
 class SecurityAuditEnv(

conftest.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# Pytest: repo root is also the `security_audit_env` package. Without this,
+# some pytest versions try to import `./__init__.py` as a test module, which
+# fails (relative imports need a parent package). Keep this file free of
+# `tests/` imports to avoid import cycles.
+
+# Ignore the package entrypoint at repo root; actual tests live in tests/ only.
+collect_ignore = ["__init__.py"]

generate_submission_plots.py

@@ -1,160 +0,0 @@
-#!/usr/bin/env python3
-"""
-Generate training plots for AISHA submission.
-Simulates training results with realistic data.
-"""
-
-import os
-import numpy as np
-import matplotlib.pyplot as plt
-from pathlib import Path
-
-# Create plots directory
-Path("./plots").mkdir(exist_ok=True)
-
-# Set random seed for reproducibility
-np.random.seed(42)
-
-# Simulate realistic training data
-episodes = np.arange(1, 11)
-
-# Baseline: random agent (low, noisy scores)
-baseline_scores = np.array([0.18, 0.22, 0.19, 0.25, 0.21, 0.23, 0.20, 0.24, 0.22, 0.26])
-
-# Pre-training: LLM agent without fine-tuning (moderate scores)
-pretrain_scores = np.array([0.32, 0.35, 0.38, 0.36, 0.40, 0.39, 0.41, 0.38, 0.42, 0.40])
-
-# Post-training: LLM agent after GRPO training (higher scores, less variance)
-posttrain_scores = np.array([0.52, 0.55, 0.58, 0.56, 0.60, 0.59, 0.61, 0.58, 0.62, 0.60])
-
-# Simulate training loss curve (decreasing over steps)
-num_steps = 150
-training_steps = np.arange(1, num_steps + 1)
-# Loss starts high and decreases with some noise
-base_loss = 2.0 * np.exp(-training_steps / 50)
-noise = np.random.normal(0, 0.05, num_steps)
-training_loss = np.maximum(base_loss + noise, 0.1)
-
-print("=" * 70)
-print("AISHA TRAINING SIMULATION - GENERATING SUBMISSION PLOTS")
-print("=" * 70)
-
-# Plot 1: Episode Rewards Comparison
-print("\n[1/3] Generating reward_per_episode.png...")
-fig, ax = plt.subplots(figsize=(12, 7))
-
-ax.plot(episodes, baseline_scores, 'r--o', label='Random Agent (Baseline)', 
-        linewidth=2.5, markersize=8, alpha=0.8)
-ax.plot(episodes, pretrain_scores, 'b-s', label='LLM Agent (Pre-training)', 
-        linewidth=2.5, markersize=8, alpha=0.8)
-ax.plot(episodes, posttrain_scores, 'g-^', label='LLM Agent (Post-training GRPO)', 
-        linewidth=2.5, markersize=8, alpha=0.8)
-
-ax.set_xlabel('Episode', fontsize=14, fontweight='bold')
-ax.set_ylabel('Total Reward (0.0 - 1.0)', fontsize=14, fontweight='bold')
-ax.set_title('AISHA: Episode Reward — Baseline vs Trained Agent', fontsize=15, fontweight='bold')
-ax.legend(fontsize=12, loc='lower right', framealpha=0.95)
-ax.grid(True, alpha=0.3, linestyle='--')
-ax.set_ylim(0, 1.0)
-ax.set_xticks(episodes)
-
-plt.tight_layout()
-plt.savefig('./plots/reward_per_episode.png', dpi=150, bbox_inches='tight')
-plt.close()
-print("   ✓ Saved: ./plots/reward_per_episode.png")
-
-# Plot 2: Training Loss Curve
-print("[2/3] Generating training_loss.png...")
-fig, ax = plt.subplots(figsize=(12, 7))
-
-ax.plot(training_steps, training_loss, 'b-', linewidth=2, alpha=0.7, label='Training Loss')
-
-# Add moving average
-window = 10
-moving_avg = np.convolve(training_loss, np.ones(window)/window, mode='valid')
-ax.plot(range(window, num_steps + 1), moving_avg, 'r-', linewidth=2.5, 
-        label=f'Moving Average (window={window})', alpha=0.9)
-
-ax.set_xlabel('Training Step', fontsize=14, fontweight='bold')
-ax.set_ylabel('Loss', fontsize=14, fontweight='bold')
-ax.set_title('AISHA: Training Loss Curve (GRPO)', fontsize=15, fontweight='bold')
-ax.legend(fontsize=12, loc='upper right', framealpha=0.95)
-ax.grid(True, alpha=0.3, linestyle='--')
-
-plt.tight_layout()
-plt.savefig('./plots/training_loss.png', dpi=150, bbox_inches='tight')
-plt.close()
-print("   ✓ Saved: ./plots/training_loss.png")
-
-# Plot 3: Performance Comparison Bar Chart
-print("[3/3] Generating performance_comparison.png...")
-fig, ax = plt.subplots(figsize=(10, 7))
-
-agents = ['Random\nAgent', 'LLM\n(Pre-train)', 'LLM\n(Post-train)']
-avgs = [
-    np.mean(baseline_scores),
-    np.mean(pretrain_scores),
-    np.mean(posttrain_scores)
-]
-colors = ['#e74c3c', '#3498db', '#2ecc71']
-stds = [
-    np.std(baseline_scores),
-    np.std(pretrain_scores),
-    np.std(posttrain_scores)
-]
-
-bars = ax.bar(agents, avgs, color=colors, width=0.6, edgecolor='black', 
-              linewidth=2, alpha=0.85, yerr=stds, capsize=10, error_kw={'linewidth': 2})
-
-ax.set_ylabel('Average Episode Reward (0.0 - 1.0)', fontsize=14, fontweight='bold')
-ax.set_title('AISHA: Average Performance Comparison', fontsize=15, fontweight='bold')
-ax.set_ylim(0, 1.0)
-ax.grid(True, alpha=0.3, axis='y', linestyle='--')
-
-# Add value labels on bars
-for bar, val, std in zip(bars, avgs, stds):
-    height = bar.get_height()
-    ax.text(bar.get_x() + bar.get_width()/2, height + std + 0.03,
-            f'{val:.3f}', ha='center', va='bottom', fontsize=13, fontweight='bold')
-
-plt.tight_layout()
-plt.savefig('./plots/performance_comparison.png', dpi=150, bbox_inches='tight')
-plt.close()
-print("   ✓ Saved: ./plots/performance_comparison.png")
-
-# Print summary
-print("\n" + "=" * 70)
-print("TRAINING RESULTS SUMMARY")
-print("=" * 70)
-print(f"\nRandom Agent (Baseline):")
-print(f"  Average Score:  {avgs[0]:.4f} ± {stds[0]:.4f}")
-print(f"  Min Score:      {np.min(baseline_scores):.4f}")
-print(f"  Max Score:      {np.max(baseline_scores):.4f}")
-
-print(f"\nLLM Agent (Pre-training):")
-print(f"  Average Score:  {avgs[1]:.4f} ± {stds[1]:.4f}")
-print(f"  Min Score:      {np.min(pretrain_scores):.4f}")
-print(f"  Max Score:      {np.max(pretrain_scores):.4f}")
-
-print(f"\nLLM Agent (Post-training GRPO):")
-print(f"  Average Score:  {avgs[2]:.4f} ± {stds[2]:.4f}")
-print(f"  Min Score:      {np.min(posttrain_scores):.4f}")
-print(f"  Max Score:      {np.max(posttrain_scores):.4f}")
-
-improvement_pretrain = ((avgs[1] - avgs[0]) / avgs[0]) * 100
-improvement_posttrain = ((avgs[2] - avgs[1]) / avgs[1]) * 100
-improvement_total = ((avgs[2] - avgs[0]) / avgs[0]) * 100
-
-print(f"\nImprovement:")
-print(f"  Pre-train vs Baseline:  {improvement_pretrain:+.1f}%")
-print(f"  Post-train vs Pre-train: {improvement_posttrain:+.1f}%")
-print(f"  Post-train vs Baseline: {improvement_total:+.1f}%")
-
-print(f"\nTraining Loss:")
-print(f"  Initial Loss:   {training_loss[0]:.4f}")
-print(f"  Final Loss:     {training_loss[-1]:.4f}")
-print(f"  Reduction:      {((training_loss[0] - training_loss[-1]) / training_loss[0]) * 100:.1f}%")
-
-print("\n" + "=" * 70)
-print("✓ All plots generated successfully!")
-print("=" * 70)

inference.py

@@ -9,21 +9,35 @@ environment. Reads API credentials from environment variables.
 ENV VARS (required):
     API_BASE_URL  — The API endpoint for the LLM
     MODEL_NAME    — The model identifier to use
-    HF_TOKEN      — Your Hugging Face / API key
+    HF_TOKEN      — API key (Hugging Face, OpenRouter, etc. — sent as client api_key)
+
+Optional:
+    INFERENCE_LOG_LLM — If set, append each raw model response to this file path
+    INFERENCE_LOG_API — If set, append API errors and fallback-to-list_tools events
+    INFERENCE_INTERACTIVE — 1 / true: wait for Enter between steps (see --interactive)
+    INFERENCE_PAUSE — step (default) | scenario: what "interactive" pauses between
+    INFERENCE_SUMMARY_FILE — If set, append a short grader + stats block per scenario
 """
 
-import json
+import argparse
 import os
-import re
 import sys
 import textwrap
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from openai import OpenAI
 
+try:
+    from dotenv import load_dotenv
+
+    load_dotenv()
+except ImportError:
+    pass
+
 # --- ENV VARS ---
-API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.3-70B-Instruct")
+# Defaults favour OpenRouter + Llama 3.2 3B; override with env or .env.
+API_BASE_URL = os.getenv("API_BASE_URL", "https://openrouter.ai/api/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/llama-3.2-3b-instruct:free")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 if HF_TOKEN is None:
@@ -71,29 +85,92 @@ CRITICAL RULES:
 """).strip()
 
 
-def parse_action(response_text: str) -> Optional[Dict[str, Any]]:
-    """Extract a JSON action from the LLM's response."""
-    if not response_text:
-        return None
+def _append_llm_log(path: str, scenario_id: str, step: int, text: str) -> None:
+    try:
+        with open(path, "a", encoding="utf-8") as f:
+            f.write(f"\n=== {scenario_id} step={step} ===\n{text}\n")
+    except OSError:
+        pass
 
-    text = response_text.strip()
-    text = re.sub(r"```json\s*", "", text)
-    text = re.sub(r"```\s*$", "", text)
-    text = text.strip()
 
+def _append_api_log(path: str, scenario_id: str, step: int, text: str) -> None:
     try:
-        return json.loads(text)
-    except json.JSONDecodeError:
+        with open(path, "a", encoding="utf-8") as f:
+            f.write(f"\n=== {scenario_id} step={step} API ===\n{text.rstrip()}\n")
+    except OSError:
         pass
 
-    match = re.search(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", text, re.DOTALL)
-    if match:
-        try:
-            return json.loads(match.group(0))
-        except json.JSONDecodeError:
-            pass
 
-    return None
+def _append_summary_file(path: str, text: str) -> None:
+    try:
+        with open(path, "a", encoding="utf-8") as f:
+            f.write(text)
+            if not text.endswith("\n"):
+                f.write("\n")
+    except OSError:
+        pass
+
+
+def _format_grader_block(scenario_id: str, grades: Dict[str, Any], episode_step_reward_sum: float) -> str:
+    """Human-readable grader output (where your final 'reward' / score comes from)."""
+    lines = [
+        "",
+        f"{'='*60}",
+        f"  REWARD / GRADER BREAKDOWN  —  scenario: {scenario_id}",
+        f"{'='*60}",
+        f"  final_score (0–1, main benchmark):  {grades.get('final_score', 0.0):.4f}",
+        f"  sum of per-step rewards (episode):  {episode_step_reward_sum:.4f}",
+        f"  true positives / total vulns:     {grades.get('true_positives', 0)}/{grades.get('total_vulnerabilities', 0)}  (detection_rate={grades.get('detection_rate', 0.0):.2f})",
+        f"  hosts examined / total hosts:        {grades.get('hosts_examined', 0)}/{grades.get('total_hosts', 0)}  (coverage={grades.get('coverage', 0.0):.2f})",
+        f"  false positives (penalty):         {grades.get('false_positives', 0)}  (fp_penalty -{grades.get('fp_penalty', 0.0):.2f})",
+        f"  severity / classification:         {grades.get('severity_accuracy', 0.0):.2f} / {grades.get('classification_accuracy', 0.0):.2f}",
+        f"  report quality:                    {grades.get('report_quality', 0.0):.2f}",
+        f"{'='*60}",
+    ]
+    return "\n".join(lines) + "\n"
+
+
+def _format_zero_score_hint(
+    n_list_tools: int,
+    n_api_errors: int,
+    total_steps: int,
+) -> str:
+    parts = [
+        "  HINT: final_score is 0 when no findings match the scenario, or coverage is near zero.",
+    ]
+    if n_list_tools >= max(1, total_steps - 1) and total_steps > 0:
+        parts.append(
+            "  → Most steps were 'list_tools' (no discovery). Use use_tool (network_scan, web_crawl) then submit_finding."
+        )
+    if n_api_errors > 0:
+        parts.append(
+            f"  → {n_api_errors} LLM API call(s) failed (see INFERENCE_LOG_API or stderr); responses may be fallbacks, not the model."
+        )
+    return "\n".join(parts) + "\n"
+
+
+def _env_bool(name: str) -> bool:
+    return os.getenv(name, "").lower() in ("1", "true", "yes", "on")
+
+
+def _wait_interactive(
+    message: str,
+) -> str:
+    """Block until the user accepts the next action. Returns a short status for logging."""
+    if not sys.stdin.isatty():
+        return "skipped (no tty)"
+    try:
+        return input(message).strip().lower() or "ok"
+    except EOFError:
+        return "eof"
+
+
+def _config_interactive() -> Tuple[bool, str]:
+    """(interactive, pause) where pause is 'step' or 'scenario'."""
+    pause = os.getenv("INFERENCE_PAUSE", "step").lower().strip()
+    if pause not in ("step", "scenario"):
+        pause = "step"
+    return _env_bool("INFERENCE_INTERACTIVE"), pause
 
 
 def build_prompt(step: int, observation: Any, history: List[str], max_steps: int = 30) -> str:
@@ -144,11 +221,28 @@ def build_prompt(step: int, observation: Any, history: List[str], max_steps: int
     return "\n".join(parts)
 
 
-def run_scenario(client: OpenAI, scenario_id: str, env_url: str) -> float:
-    """Run the agent on one scenario and return the final score."""
-    from security_audit_env import SecurityAuditEnv, SecurityAuditAction
+def run_scenario(
+    client: OpenAI,
+    scenario_id: str,
+    env_url: str,
+    *,
+    interactive: bool = False,
+    pause: str = "step",
+) -> float:
+    """Run the agent on one scenario and return the final score.
+
+    If ``interactive`` and ``pause == "step"``, wait for Enter after each step
+    (before the next LLM call) to space out API traffic and avoid rate limits.
+    If ``pause == "scenario"``, only :func:`main` pauses between scenarios.
+    """
+    from security_audit_env import (
+        SecurityAuditAction,
+        SecurityAuditEnv,
+        parse_llm_action_text,
+    )
 
     max_steps = SCENARIO_MAX_STEPS.get(scenario_id, 30)
+    api_log = os.getenv("INFERENCE_LOG_API")
 
     print(f"\n{'='*60}")
     print(f"Running scenario: {scenario_id} (max {max_steps} steps)")
@@ -162,13 +256,50 @@ def run_scenario(client: OpenAI, scenario_id: str, env_url: str) -> float:
     total_steps = 0
     success = False
     last_error = None
+    user_quit_scenario = False
+    last_grades: Optional[Dict[str, Any]] = None
+    n_list_tools = 0
+    n_api_errors = 0
+    summary_path = os.getenv("INFERENCE_SUMMARY_FILE")
 
     try:
         with SecurityAuditEnv(base_url=env_url).sync() as env:
+            if interactive and pause == "step" and sys.stdin.isatty():
+                u = _wait_interactive(
+                    f"\n>>> Starting '{scenario_id}'. Press Enter to run the first step (LLM call), or 'q' + Enter to skip this scenario.\n> "
+                )
+                if u == "q":
+                    return 0.0
+
             result = env.reset(scenario_id=scenario_id)
             observation = result.observation
             history: List[str] = []
 
+            def _do_force_report() -> None:
+                nonlocal result, all_rewards, total_steps, final_score, success, last_error, observation, last_grades
+                try:
+                    act = SecurityAuditAction(action_type="generate_report")
+                    result = env.step(act)
+                    reward = result.reward or 0.0
+                    all_rewards.append(reward)
+                    total_steps = total_steps + 1
+                    _ts = total_steps
+                    _cum = sum(all_rewards)
+                    print(
+                        f"[STEP]  step={_ts} action=generate_report reward={reward:.2f} "
+                        f"cum={_cum:.2f} done={str(result.done).lower()} error=null",
+                        flush=True,
+                    )
+                    observation = result.observation
+                    grades = getattr(observation, "metadata", {}) or {}
+                    grades = grades.get("grades", {})
+                    last_grades = grades if isinstance(grades, dict) and grades else None
+                    final_score = grades.get("final_score", reward) if last_grades else (reward or 0.0)
+                    success = final_score > 0
+                except Exception as exc:
+                    final_score = 0.0
+                    last_error = str(exc)
+
             for step in range(1, max_steps + 1):
                 if result.done:
                     break
@@ -190,28 +321,34 @@ def run_scenario(client: OpenAI, scenario_id: str, env_url: str) -> float:
                     )
                     response_text = completion.choices[0].message.content or ""
                 except Exception as exc:
+                    n_api_errors += 1
                     last_error = str(exc)
                     response_text = '{"action_type": "list_tools"}'
+                    err_line = f"[API error — using fallback list_tools] {type(exc).__name__}: {exc}"
+                    if api_log:
+                        _append_api_log(api_log, scenario_id, step, err_line)
+                    else:
+                        print(f"  {err_line}", flush=True)
+
+                log_path = os.getenv("INFERENCE_LOG_LLM")
+                if log_path and response_text:
+                    _append_llm_log(log_path, scenario_id, step, response_text)
+
+                llm_action, json_err = parse_llm_action_text(response_text)
+                if llm_action is None:
+                    last_error = json_err or "Could not parse LLM action JSON"
+                    action = SecurityAuditAction(action_type="list_tools")
+                else:
+                    last_error = None
+                    action = llm_action.to_security_audit_action()
+                if action.action_type == "list_tools":
+                    n_list_tools += 1
 
-                action_dict = parse_action(response_text)
-                if not action_dict:
-                    last_error = "Could not parse LLM response as JSON"
-                    action_dict = {"action_type": "list_tools"}
-
-                action_type = action_dict.get("action_type", "list_tools")
-                tool_name = action_dict.get("tool_name")
-                arguments = action_dict.get("arguments", {})
-
-                action_str = action_type
-                if tool_name:
-                    action_str += f"({tool_name})"
+                action_str = action.action_type
+                if action.tool_name:
+                    action_str += f"({action.tool_name})"
 
                 try:
-                    action = SecurityAuditAction(
-                        action_type=action_type,
-                        tool_name=tool_name,
-                        arguments=arguments,
-                    )
                     result = env.step(action)
                     observation = result.observation
                     last_error = None
@@ -222,48 +359,72 @@ def run_scenario(client: OpenAI, scenario_id: str, env_url: str) -> float:
                     total_steps = step
                     # --- MANDATORY STDOUT: [STEP] ---
                     error_str = last_error.replace("\n", " ") if last_error else "null"
-                    print(f"[STEP]  step={step} action={action_str} reward={reward:.2f} done=false error={error_str}", flush=True)
+                    _c = sum(all_rewards)
+                    print(
+                        f"[STEP]  step={step} action={action_str} reward={reward:.2f} "
+                        f"cum={_c:.2f} done=false error={error_str}",
+                        flush=True,
+                    )
                     break
 
                 reward = result.reward or 0.0
                 all_rewards.append(reward)
                 total_steps = step
+                _cum = sum(all_rewards)
 
                 history.append(f"Step {step}: {action_str} → reward {reward:+.2f}")
 
                 # --- MANDATORY STDOUT: [STEP] ---
                 done_str = "true" if result.done else "false"
                 error_str = last_error.replace("\n", " ") if last_error else "null"
-                print(f"[STEP]  step={step} action={action_str} reward={reward:.2f} done={done_str} error={error_str}", flush=True)
+                print(
+                    f"[STEP]  step={step} action={action_str} reward={reward:.2f} "
+                    f"cum={_cum:.2f} done={done_str} error={error_str}",
+                    flush=True,
+                )
 
                 if result.done:
                     grades = getattr(observation, "metadata", {}) or {}
                     grades = grades.get("grades", {})
-                    final_score = grades.get("final_score", reward)
+                    last_grades = grades if isinstance(grades, dict) and grades else None
+                    final_score = grades.get("final_score", reward) if last_grades else 0.0
                     success = final_score > 0
                     break
-            else:
-                # Didn't finish — force report generation
-                try:
-                    action = SecurityAuditAction(action_type="generate_report")
-                    result = env.step(action)
-                    reward = result.reward or 0.0
-                    all_rewards.append(reward)
-                    total_steps += 1
 
-                    done_str = "true" if result.done else "false"
-                    print(f"[STEP]  step={total_steps} action=generate_report reward={reward:.2f} done={done_str} error=null", flush=True)
-
-                    grades = getattr(result.observation, "metadata", {}) or {}
-                    grades = grades.get("grades", {})
-                    final_score = grades.get("final_score", 0.0)
-                    success = final_score > 0
-                except Exception as exc:
-                    final_score = 0.0
-                    last_error = str(exc)
+                if interactive and pause == "step" and sys.stdin.isatty() and not result.done:
+                    u2 = _wait_interactive(
+                        f"\n>>> {scenario_id} step {step}/{max_steps} done. "
+                        "Press Enter for the next LLM call, or 'q' + Enter to end this scenario (a report will be generated).\n> "
+                    )
+                    if u2 == "q":
+                        user_quit_scenario = True
+                        break
+            else:
+                # No break — ran all steps without terminal done: force report
+                _do_force_report()
+            if user_quit_scenario:
+                _do_force_report()
     except Exception as exc:
         last_error = str(exc)
     finally:
+        if last_grades is not None:
+            _sm = _format_grader_block(scenario_id, last_grades, sum(all_rewards))
+            print(_sm, flush=True)
+            if summary_path:
+                _append_summary_file(summary_path, _sm)
+        elif total_steps > 0:
+            _mini = (
+                f"\n  (No grader report in metadata — score may be unset. "
+                f"Steps={total_steps} list_tools_steps≈{n_list_tools} api_errors={n_api_errors})\n"
+            )
+            print(_mini, flush=True)
+            if summary_path:
+                _append_summary_file(summary_path, _mini)
+        if final_score == 0.0 and (last_grades is not None or total_steps > 0):
+            _hint = _format_zero_score_hint(n_list_tools, n_api_errors, total_steps)
+            print(_hint, flush=True)
+            if summary_path:
+                _append_summary_file(summary_path, _hint)
         # --- MANDATORY STDOUT: [END] (always emitted, even on exception) ---
         rewards_str = ",".join(f"{r:.2f}" for r in all_rewards)
         success_str = "true" if success else "false"
@@ -272,19 +433,68 @@ def run_scenario(client: OpenAI, scenario_id: str, env_url: str) -> float:
     return final_score
 
 
-def main():
+def _parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Run the baseline LLM agent on SecurityAuditEnv.",
+    )
+    p.add_argument(
+        "-i",
+        "--interactive",
+        action="store_true",
+        help="Wait for your input between LLM steps (or between scenarios) to space out API calls and reduce rate limits",
+    )
+    p.add_argument(
+        "--pause",
+        choices=["step", "scenario"],
+        default=None,
+        help="With --interactive: 'step' pauses after each environment step; 'scenario' only between easy/medium/hard",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
     """Run baseline inference across all scenarios."""
+    args = _parse_args()
+    env_inter = args.interactive or _env_bool("INFERENCE_INTERACTIVE")
+    pause = args.pause or os.getenv("INFERENCE_PAUSE", "step")
+    if pause not in ("step", "scenario"):
+        pause = "step"
+
     print("Security Audit Environment — Baseline Inference")
+    if env_inter:
+        print("Mode: INTERACTIVE (you control the pace; stdin must be a TTY)")
+        print(f"  Pause: {pause}  (INFERENCE_PAUSE, or --pause)")
     print(f"API: {API_BASE_URL}")
     print(f"Model: {MODEL_NAME}")
 
     llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
     env_url = os.getenv("ENV_URL", "http://localhost:8000")
 
-    scores = {}
-    for scenario_id in SCENARIOS:
+    scores: Dict[str, float] = {}
+    for i, scenario_id in enumerate(SCENARIOS):
+        if env_inter and pause == "scenario" and sys.stdin.isatty():
+            if i == 0:
+                nxt0 = _wait_interactive(
+                    f"\n>>> Press Enter to start the first scenario ('{scenario_id}'), or 'q' + Enter to cancel.\n> "
+                )
+                if nxt0 == "q":
+                    print("(Cancelled.)", flush=True)
+                    return
+            else:
+                nxt = _wait_interactive(
+                    f"\n>>> Previous scenario(s) finished. Press Enter to start '{scenario_id}', or 'q' + Enter to stop the run.\n> "
+                )
+                if nxt == "q":
+                    print("(Stopping — remaining scenarios skipped.)", flush=True)
+                    break
         try:
-            score = run_scenario(llm_client, scenario_id, env_url)
+            score = run_scenario(
+                llm_client,
+                scenario_id,
+                env_url,
+                interactive=env_inter,
+                pause=pause,
+            )
             scores[scenario_id] = score
         except Exception as exc:
             print(f"  ERROR on {scenario_id}: {exc}")
@@ -293,9 +503,11 @@ def main():
     print(f"\n{'='*60}")
     print("BASELINE SCORES")
     print(f"{'='*60}")
-    for sid, score in scores.items():
-        print(f"  {sid:10s}: {score:.4f}")
-    avg = sum(scores.values()) / len(scores) if scores else 0.0
+    for sid in SCENARIOS:
+        if sid in scores:
+            print(f"  {sid:10s}: {scores[sid]:.4f}")
+    rans = [scores[k] for k in SCENARIOS if k in scores]
+    avg = sum(rans) / len(rans) if rans else 0.0
     print(f"  {'average':10s}: {avg:.4f}")
     print(f"{'='*60}")
 

models.py

@@ -10,10 +10,12 @@ Simulates real-world VAPT (Vulnerability Assessment & Penetration Testing)
 engagements where an AI agent audits infrastructure for security compliance.
 """
 
-from typing import Any, Dict, List, Literal, Optional
+import json
+import re
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 from openenv.core.env_server.types import Action, Observation, State
-from pydantic import Field
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 
 class SecurityAuditAction(Action):
@@ -41,6 +43,79 @@ class SecurityAuditAction(Action):
     )
 
 
+class LLMJsonAction(BaseModel):
+    """Wire JSON for one model turn, validated before ``SecurityAuditAction``.
+
+    Unknown top-level keys are ignored so minor format drift does not fail parsing.
+    """
+
+    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)
+
+    action_type: Literal["list_tools", "use_tool", "submit_finding", "generate_report"] = Field(
+        ...,
+        description="Which environment action to take",
+    )
+    tool_name: Optional[str] = Field(
+        default=None,
+        description="Tool name when action_type is use_tool",
+    )
+    arguments: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments for use_tool or fields for submit_finding",
+    )
+
+    def to_security_audit_action(self) -> SecurityAuditAction:
+        return SecurityAuditAction(
+            action_type=self.action_type,
+            tool_name=self.tool_name,
+            arguments=self.arguments,
+        )
+
+
+def extract_json_object_from_text(raw: str) -> Optional[Dict[str, Any]]:
+    """Return the first JSON object from model text, or None."""
+    if not (raw and raw.strip()):
+        return None
+
+    text = raw.strip()
+    text = re.sub(r"```json\s*", "", text)
+    text = re.sub(r"```\s*$", "", text, flags=re.MULTILINE)
+    text = text.strip()
+
+    try:
+        val = json.loads(text)
+    except json.JSONDecodeError:
+        val = None
+
+    if isinstance(val, dict):
+        return val
+
+    match = re.search(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", text, re.DOTALL)
+    if match:
+        try:
+            v2 = json.loads(match.group(0))
+        except json.JSONDecodeError:
+            return None
+        if isinstance(v2, dict):
+            return v2
+
+    return None
+
+
+def parse_llm_action_text(raw: str) -> Tuple[Optional[LLMJsonAction], Optional[str]]:
+    """Parse and validate one action from a chat message.
+
+    Returns (model, None) on success, or (None, error_message) on failure.
+    """
+    data = extract_json_object_from_text(raw)
+    if data is None:
+        return None, "Could not extract a JSON object from model response"
+    try:
+        return LLMJsonAction.model_validate(data), None
+    except ValidationError as exc:
+        return None, str(exc)
+
+
 class SecurityAuditObservation(Observation):
     """Observation returned after each step.
 

notebook_builder.py

@@ -1 +0,0 @@
-# Placeholder

pyproject.toml

@@ -8,6 +8,11 @@
 requires = ["setuptools>=45", "wheel"]
 build-backend = "setuptools.build_meta"
 
+# Only run tests from tests/ — avoids collecting repo-root __init__.py (package) as a module.
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]
+
 [project]
 name = "openenv-security_audit_env"
 version = "0.1.0"
@@ -19,6 +24,7 @@ dependencies = [
     # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
     "openenv-core[core]>=0.2.3",
     "openai>=1.0.0",
+    "python-dotenv>=1.0.0",
 ]
 
 [project.optional-dependencies]