Transforming Agent Demo Code into a Reusable Python Package

The shift from experimental code to a production-ready system often begins with a simple question: Why can’t we reuse this?

In the latest installment of the Agent Series, developers faced this exact scenario with a 900-line Python script designed to showcase eight security layers for AI agents. While effective as a demonstration tool, the monolithic harness_full_demo.py file suffered from critical limitations: tightly coupled components that couldn’t be tested in isolation, no importable API, and a design that prioritized illustration over practicality. The solution? A complete architectural overhaul to create harness, a modular Python package that transforms fragile demo code into a robust, reusable framework for production environments.

Designing for Reusability: A Modular Architecture

The new package adopts a layered approach, separating concerns to improve maintainability and security. Each module addresses a specific function while exposing a clean public API through harness/__init__.py. The structure reflects this philosophy:

harness/
├── __init__.py          # Public API exports
├── registry.py          # Action registration and permission levels
├── budget.py            # Permission budget management with refund capability
├── sandbox.py           # Input sanitization and secure evaluation
├── audit.py             # Immutable audit logging with hash-chained verification
├── rollback.py          # Transaction rollback for state consistency
└── harness.py           # Unified entry point: AgentHarness

This separation enables components to be tested independently, imported into other projects, and extended without disrupting existing functionality.

Key API Decisions: Clarity and Safety

Three critical design choices underpin the package’s functionality:

1. Permission Levels and Explicit Error Handling

The registry.py module introduces PermissionLevel, an enumeration defining four tiers of access: READ, WRITE, ADMIN, and IRREVERSIBLE. Each action registered with ActionRegistry must specify a level and associated budget cost. The design deliberately avoids Python’s __getitem__ magic method in favor of a get() method that raises a consistent PermissionError—omitting internal KeyError details to prevent information leakage. This approach simplifies error handling for calling code while maintaining security boundaries.

2. Refundable Budget Management

A major flaw in the original implementation was irreversible budget depletion. The PermissionBudget class in budget.py now includes a refund() method, allowing budgets to be restored when actions are rejected—particularly critical for IRREVERSIBLE operations that trigger human approval. This change ensures financial accounting remains accurate, even when automated flows are interrupted by safety checks.

3. Separate Execution Paths for Human and Machine Workflows

The AgentHarness class exposes two distinct methods: execute() for automated paths and approve_and_execute() for human-approved actions. Combining these into a single method with a boolean flag would obscure intent and complicate testing. By keeping them separate, the API remains unambiguous: execute() assumes pre-approval, while approve_and_execute() requires explicit confirmation.

Security Enhancements: Sandboxing and Audit Trails

Security isn’t an afterthought in this package—it’s woven into the core design.

Input Sanitization with Comprehensive Pattern Matching

The sandbox.py module uses a robust regular expression to detect and block injection attempts:

INJECTION_PATTERN = re.compile(
    r"(ignore.*(previous|above|prior)|forget.*instruction|"
    r"you are now|act as|jailbreak|bypass|"
    r"override.*system|system.*override|"
    r"</s>|\n\n###|###\s*system|<\|im_start\||>|system prompt)",
    re.IGNORECASE
)

This pattern covers both SYSTEM OVERRIDE and override system variations, addressing bugs discovered during adversarial testing in prior iterations. The inclusion of \n\n### ensures newline-separated inputs are properly sanitized, closing a critical gap in the original implementation.

Immutable Audit Logging

The audit.py module introduces ImmutableAuditLog, a class that records every action with cryptographic integrity. Each log entry includes a prev_hash field, creating a hash chain that detects tampering. The verify_integrity() method replays the chain to confirm no fields have been altered, while __len__() provides a simple way to count entries during testing. This design ensures a tamper-evident record of all operations, critical for compliance and debugging.

Transaction Safety: Rollback and State Management

The rollback.py module adds transactional safety to stateful operations. Its transaction() context manager captures a deep copy of the state before execution:

@contextmanager
def transaction(self, state: dict, op_name: str):
    snapshot = copy.deepcopy(state)
    self._snapshots.append({"op": op_name, "snapshot": snapshot})
    try:
        yield state
    except Exception:
        state.clear()
        state.update(snapshot)
        self._snapshots.pop()
        raise

Manual rollbacks are also supported via rollback_last(), which restores the most recent snapshot. This feature is invaluable when automated workflows encounter unexpected failures, enabling recovery without data loss.

Putting It All Together: Real-World Integration

The refactored package supports two primary integration styles:

Standalone Python Applications

A basic workflow demonstrates how components interact:

harness = AgentHarness(budget=50)

# Register actions with distinct permission levels and costs
harness.registry.register(
    RegisteredAction(
        "read_ticket",
        PermissionLevel.READ,
        1,
        "Read Jira ticket",
        handler_fn
    )
)
harness.registry.register(
    RegisteredAction(
        "write_draft",
        PermissionLevel.WRITE,
        3,
        "Write draft fix",
        handler_fn
    )
)

Executing a sequence of permitted actions deducts from the budget:

# Total cost: 1 (read) + 3 (write) + 8 (admin) = 12
r1 = harness.execute("read_ticket", ticket_id="BUG-101")
r2 = harness.execute("write_draft", ticket_id="BUG-101", patch="fix: add null check")
r3 = harness.execute("create_pr", ticket_id="BUG-101", title="fix: BUG-101")  # 38 remaining

Graph-Based Workflows with LangGraph

For more complex systems, the package can be embedded into graph-based frameworks like LangGraph, enabling orchestration of multi-step agentic processes while maintaining centralized control over permissions and budgeting.

Looking Ahead: Production-Ready AI Security

The evolution from a 900-line demo file to a modular package marks a significant milestone in the Agent Series. By addressing design flaws, enhancing security, and prioritizing reusability, developers have created a framework that bridges the gap between experimental code and enterprise-grade AI systems. Future iterations may explore dynamic permission scaling, integration with third-party identity providers, or enhanced sandboxing techniques—each building on the solid foundation established here.