Control plane — permissions, plan mode & blocking hooks

from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent.llms.base import LLMResponse
from shipit_agent.models import ToolCall


class ScriptedLLM:
    """A deterministic, offline LLM: replays a fixed list of responses."""

    def __init__(self, responses):
        self._responses = list(responses)
        self._i = 0

    def complete(self, *, messages, tools=None, system_prompt=None,
                 metadata=None, **kwargs):
        # Return the next scripted response; after the script is exhausted,
        # always finish so a denied/looping tool can't hang the run.
        if self._i < len(self._responses):
            resp = self._responses[self._i]
            self._i += 1
            return resp
        return LLMResponse(content="done")


def tool_call_then_done(tool_name, arguments):
    """Script: call one tool, then return a final answer."""
    return [
        LLMResponse(tool_calls=[ToolCall(name=tool_name, arguments=arguments)]),
        LLMResponse(content="All done."),
    ]

from shipit_agent import Agent, FunctionTool


def bash(command: str = "") -> str:
    return f"(pretend) ran: {command}"


def file_delete(path: str = "") -> str:
    return f"(pretend) deleted: {path}"


def sql(query: str = "") -> str:
    return f"(pretend) executed SQL: {query}"


def read_file(path: str = "") -> str:
    return f"(pretend) contents of {path}"


def write_file(path: str = "", text: str = "") -> str:
    return f"(pretend) wrote {len(text)} bytes to {path}"


STUB_TOOLS = [
    FunctionTool.from_callable(bash, name="bash", description="Run a shell command."),
    FunctionTool.from_callable(file_delete, name="file_delete", description="Delete a file."),
    FunctionTool.from_callable(sql, name="sql", description="Execute SQL."),
    FunctionTool.from_callable(read_file, name="read_file", description="Read a file."),
    FunctionTool.from_callable(write_file, name="write_file", description="Write a file."),
]
[t.name for t in STUB_TOOLS]

from shipit_agent import PermissionEngine

engine = PermissionEngine(
    deny=["bash", "*_delete"],   # bash and anything ending in _delete are blocked
    ask=["sql"],                 # sql needs approval
    allow=["read*"],             # read_file is auto-allowed
)

# The model tries to call the denied `bash` tool.
agent = Agent(
    llm=ScriptedLLM(tool_call_then_done("bash", {"command": "rm -rf /"})),
    tools=STUB_TOOLS,
    permissions=engine,
)
result = agent.run("Clean up the workspace.")

denied = [e for e in result.events if e.type == "tool_denied"]
print("final output:", result.output)
print("tool_denied events:", [(e.message, e.payload.get("reason")) for e in denied])

denied_msgs = [
    m for m in result.messages
    if getattr(m, "role", None) == "tool" and "was NOT run" in (m.content or "")
]
for m in denied_msgs:
    print(m.name, "->", m.content)

from shipit_agent import PermissionDecision

for name in ["bash", "file_delete", "sql", "read_file", "render_chart"]:
    res = engine.check(name, {})
    print(f"{name:14} -> {res.decision.value:6}  ({res.reason})")

plan_agent = Agent(
    llm=ScriptedLLM([
        LLMResponse(tool_calls=[ToolCall(name="write_file",
                                         arguments={"path": "out.txt", "text": "hi"})]),
        LLMResponse(content="Plan: 1) read inputs  2) write out.txt  3) verify."),
    ]),
    tools=STUB_TOOLS,
)
plan_result = plan_agent.plan("Generate out.txt from the inputs.")

print("output:", plan_result.output)
print("denied in plan mode:",
      [(e.message, e.payload.get("decision")) for e in plan_result.events
       if e.type == "tool_denied"])

from shipit_agent import AgentHooks

hooks = AgentHooks()


@hooks.on_before_tool
def block_destructive(name, args):
    if name == "bash" and "rm -rf" in args.get("command", ""):
        return {"decision": "deny", "reason": "destructive command blocked by hook"}
    return None  # observe-only otherwise


hooked = Agent(
    llm=ScriptedLLM(tool_call_then_done("bash", {"command": "rm -rf /tmp/data"})),
    tools=STUB_TOOLS,
    hooks=hooks,
)
hr = hooked.run("Wipe /tmp/data")
print([(e.message, e.payload.get("reason")) for e in hr.events if e.type == "tool_denied"])

from shipit_agent import PermissionResult, PermissionDecision

rewrite_hooks = AgentHooks()


@rewrite_hooks.on_before_tool
def scope_path(name, args):
    if name == "write_file":
        safe = {**args, "path": f"sandbox/{args.get('path', 'out.txt')}"}
        return PermissionResult(PermissionDecision.ALLOW, updated_arguments=safe)
    return None


rewrite_agent = Agent(
    llm=ScriptedLLM(tool_call_then_done("write_file", {"path": "secret.txt", "text": "x"})),
    tools=STUB_TOOLS,
    hooks=rewrite_hooks,
)
rr = rewrite_agent.run("Write secret.txt")
# The tool message reflects the rewritten path (sandbox/secret.txt).
for m in rr.messages:
    if getattr(m, "role", None) == "tool" and m.name == "write_file":
        print(m.content)

redact_hooks = AgentHooks()


@redact_hooks.on_user_prompt
def redact(prompt):
    return prompt.replace("SECRET", "[redacted]")


redact_agent = Agent(llm=ScriptedLLM([LLMResponse(content="ack")]), hooks=redact_hooks)
res = redact_agent.run("My password is SECRET, keep it safe.")
# The redacted prompt is what landed in the conversation history.
user_msgs = [m for m in res.messages if getattr(m, "role", None) == "user"]
print(user_msgs[-1].content)

def approver(name, args):
    # Auto-approve reads; deny anything that touches production.
    if name.startswith("read"):
        return PermissionResult(PermissionDecision.ALLOW, reason="reads are safe")
    if "prod" in str(args.get("path", "")):
        return PermissionResult(PermissionDecision.DENY, reason="prod is off-limits")
    return None


cb_agent = Agent(
    llm=ScriptedLLM(tool_call_then_done("write_file", {"path": "prod/db", "text": "x"})),
    tools=STUB_TOOLS,
    permission_callback=approver,
)
cb = cb_agent.run("Write to prod/db")
print([(e.message, e.payload.get("reason")) for e in cb.events if e.type == "tool_denied"])