PhoneWork/docs/claude/test_can_use_tool_ask.py

"""Example: can_use_tool callback — intercept AskUserQuestion and provide answers.

Demonstrates:
- can_use_tool callback for tool permission control
- Detecting AskUserQuestion tool calls
- Using PermissionResultAllow(updated_input=...) to pre-fill user answers
- Auto-allowing read-only tools

This pattern is used by the secretary model to forward CC questions
to Feishu users and relay their responses back.

Usage:
    cd docs/claude
    ../../.venv/Scripts/python test_can_use_tool_ask.py
"""

from __future__ import annotations

import asyncio
import json

from _sdk_test_common import auth_env, make_tmpdir, remove_tmpdir, setup_auth

# Simulated user responses (in production, these come from Feishu card callbacks)
SIMULATED_ANSWERS: dict[str, str] = {}
permission_log: list[dict] = []


async def permission_callback(tool_name, input_data, context):
    """can_use_tool callback that intercepts AskUserQuestion.

    For AskUserQuestion:
      - Extracts questions and options from input_data
      - Provides pre-filled answers via updated_input
      - Returns PermissionResultAllow with the modified input

    For other tools:
      - Read-only tools: auto-allow
      - Write tools: auto-allow (for testing)
    """
    from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny

    permission_log.append({
        "tool": tool_name,
        "input_keys": list(input_data.keys()),
    })

    if tool_name == "AskUserQuestion":
        questions = input_data.get("questions", [])
        print(f"\n  🔔 AskUserQuestion intercepted! ({len(questions)} questions)")

        # Build answers dict: {question_text: selected_option_label}
        answers = {}
        for q in questions:
            question_text = q.get("question", "")
            options = q.get("options", [])
            multi = q.get("multiSelect", False)

            print(f"     Q: {question_text}")
            for opt in options:
                print(f"        - {opt['label']}: {opt.get('description', '')[:60]}")

            # In production: send card to Feishu, wait for user selection
            # Here: use first option as simulated answer
            if options:
                selected = SIMULATED_ANSWERS.get(question_text, options[0]["label"])
                answers[question_text] = selected
                print(f"     → Selected: {selected}")

        # Pre-fill answers in the tool input via updated_input
        modified_input = dict(input_data)
        if "answers" not in modified_input:
            modified_input["answers"] = {}
        if isinstance(modified_input["answers"], dict):
            modified_input["answers"].update(answers)

        print(f"     → updated_input.answers = {answers}")
        return PermissionResultAllow(updated_input=modified_input)

    # Auto-allow everything else for this test
    return PermissionResultAllow()


async def main() -> int:
    ok, msg = setup_auth()
    print(msg)
    if not ok:
        return 1

    from claude_agent_sdk import (
        AssistantMessage,
        ClaudeAgentOptions,
        ClaudeSDKClient,
        ResultMessage,
        SystemMessage,
        TextBlock,
        ToolUseBlock,
    )

    tmpdir = make_tmpdir("ask-test-")
    try:
        print("--- Test: can_use_tool intercepts AskUserQuestion ---")
        print(f"    tmpdir: {tmpdir}")

        opts = ClaudeAgentOptions(
            cwd=str(tmpdir),
            permission_mode="default",
            can_use_tool=permission_callback,
            max_turns=5,
            env=auth_env(),
        )

        # Ask Claude to ask the user a question — this will trigger AskUserQuestion
        prompt = (
            "I need you to ask the user a question using the AskUserQuestion tool. "
            "Ask them: 'Which programming language do you prefer?' "
            "with options: 'Python' (great for AI), 'TypeScript' (great for web). "
            "Then tell me what they chose."
        )

        print(f"\n  Prompt: {prompt[:100]}...")

        async with ClaudeSDKClient(opts) as client:
            await client.query(prompt)

            ask_seen = False
            result_text = ""

            async for msg in client.receive_response():
                if isinstance(msg, AssistantMessage):
                    for block in msg.content:
                        if isinstance(block, TextBlock):
                            print(f"  Claude: {block.text[:200]}")
                        elif isinstance(block, ToolUseBlock):
                            if block.name == "AskUserQuestion":
                                ask_seen = True
                                print(f"  [ToolUse] AskUserQuestion called!")
                            else:
                                print(f"  [ToolUse] {block.name}")
                elif isinstance(msg, ResultMessage):
                    result_text = msg.result or ""
                    print(f"\n  Result: {result_text[:300]}")

        print(f"\n  Permission log ({len(permission_log)} entries):")
        for entry in permission_log:
            print(f"    {entry['tool']}: keys={entry['input_keys']}")

        if ask_seen:
            print("\n  ✅ PASS: AskUserQuestion was intercepted by can_use_tool")
        else:
            print("\n  ⚠️  AskUserQuestion was not called (Claude may have answered directly)")
            print("     This is OK — it means Claude didn't need to ask.")

        return 0
    finally:
        remove_tmpdir(tmpdir)


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))