S04 Hooks：把扩展逻辑挂在循环上

HOOKS = {
    "UserPromptSubmit": [],
    "PreToolUse": [],
    "PostToolUse": [],
    "Stop": [],
}

def register_hook(event, callback):
    HOOKS[event].append(callback)

def trigger_hooks(event, *args):
    for callback in HOOKS[event]:
        result = callback(*args)
        if result is not None:
            return result
    return None

def context_inject_hook(query: str):
    print(f"[HOOK] working in {WORKDIR}")
    return None

register_hook("UserPromptSubmit", context_inject_hook)

trigger_hooks("UserPromptSubmit", query)
history.append({"role": "user", "content": query})
agent_loop(history)

def permission_hook(block):
    if block.name == "bash":
        command = block.input.get("command", "")
        if any(pattern in command for pattern in DENY_LIST):
            return "Permission denied by deny list"
        if any(word in command for word in DESTRUCTIVE):
            choice = input("Allow? [y/N] ").strip().lower()
            if choice not in ("y", "yes"):
                return "Permission denied by user"

    if block.name in ("write_file", "edit_file"):
        path = block.input.get("path", "")
        target = (WORKDIR / path).resolve()
        if not target.is_relative_to(WORKDIR):
            choice = input("Writing outside workspace. Allow? [y/N] ").strip().lower()
            if choice not in ("y", "yes"):
                return "Permission denied by user"
    return None

def log_hook(block):
    print(f"[HOOK] {block.name}")
    return None

register_hook("PreToolUse", permission_hook)
register_hook("PreToolUse", log_hook)

blocked = trigger_hooks("PreToolUse", block)
if blocked:
    results.append({
        "type": "tool_result",
        "tool_use_id": block.id,
        "content": str(blocked),
    })
    continue

output = TOOL_HANDLERS[block.name](**block.input)
trigger_hooks("PostToolUse", block, output)

def large_output_hook(block, output):
    if len(str(output)) > 100000:
        print(f"[HOOK] Large output from {block.name}")
    return None

register_hook("PostToolUse", large_output_hook)

if response.stop_reason != "tool_use":
    force = trigger_hooks("Stop", messages)
    if force:
        messages.append({"role": "user", "content": force})
        continue
    return

def summary_hook(messages):
    tool_count = sum(
        1 for message in messages
        for block in (message.get("content")
                      if isinstance(message.get("content"), list) else [])
        if isinstance(block, dict) and block.get("type") == "tool_result"
    )
    print(f"Session used {tool_count} tool calls")
    return None

register_hook("Stop", summary_hook)

trigger_hooks("UserPromptSubmit", query)
messages.append(user(query))

while True:
    response = call_model(system=SYSTEM, tools=TOOLS, messages=messages)
    messages.append(assistant(response.content))
    calls = get_tool_use_blocks(response.content)

    if not calls:
        follow_up = trigger_hooks("Stop", messages)
        if follow_up:
            messages.append(user(follow_up))
            continue
        break

    results = []
    for call in calls:
        blocked = trigger_hooks("PreToolUse", call)
        if blocked:
            output = str(blocked)
        else:
            output = TOOL_HANDLERS[call.name](**call.input)
            trigger_hooks("PostToolUse", call, output)
        results.append(tool_result(call.id, output))
    messages.append(user(results))