Keep stdin open for can_use_tool and allow it with string prompts

e2e-tests/test_tool_permissions.py

@@ -10,7 +10,9 @@
     ClaudeSDKClient,
     PermissionResultAllow,
     PermissionResultDeny,
+    ResultMessage,
     ToolPermissionContext,
+    query,
 )
 
 
@@ -63,3 +65,79 @@ async def permission_callback(
         # Clean up
         if test_path.exists():
             test_path.unlink()
+
+
+async def _run_query_with_only_can_use_tool(prompt_factory):
+    """Drive one-shot query() with can_use_tool as the *only* control-protocol
+    consumer (no hooks, no SDK MCP servers) and return the callback log."""
+    callback_invocations: list[str] = []
+    unique_id = uuid.uuid4().hex[:8]
+    # Outside cwd and /tmp so no sandbox or allow-rule can auto-approve it: the
+    # Write must round-trip through the permission callback.
+    test_path = Path.home() / f".agent_sdk_permission_e2e_{unique_id}.txt"
+
+    async def permission_callback(
+        tool_name: str,
+        input_data: dict,
+        context: ToolPermissionContext,
+    ) -> PermissionResultAllow | PermissionResultDeny:
+        callback_invocations.append(tool_name)
+        return PermissionResultAllow()
+
+    options = ClaudeAgentOptions(
+        can_use_tool=permission_callback,
+        disallowed_tools=["Bash"],
+        max_turns=3,
+    )
+    prompt_text = (
+        f"Use the Write tool to create the file {test_path} containing exactly "
+        "the text hello. Do not use any other tool."
+    )
+
+    try:
+        result = None
+        async for message in query(prompt=prompt_factory(prompt_text), options=options):
+            if isinstance(message, ResultMessage):
+                result = message
+        assert result is not None, "query() should yield a ResultMessage"
+        assert not result.is_error, f"run should succeed, got: {result}"
+        assert "Write" in callback_invocations, (
+            "can_use_tool should have been invoked for Write, "
+            f"got: {callback_invocations}"
+        )
+        assert test_path.exists(), "the permitted Write should have run"
+    finally:
+        if test_path.exists():
+            test_path.unlink()
+
+
+@pytest.mark.e2e
+@pytest.mark.anyio
+async def test_query_async_iterable_prompt_with_only_can_use_tool():
+    """query() with a finite AsyncIterable prompt and only can_use_tool must
+    keep stdin open long enough for the permission verdict to reach the CLI.
+
+    Regression: stdin used to close as soon as the prompt iterable finished
+    unless hooks or SDK MCP servers were also configured, so the CLI's
+    permission request failed with "Stream closed" and the callback never ran.
+    """
+
+    def make_stream(text: str):
+        async def stream():
+            yield {
+                "type": "user",
+                "message": {"role": "user", "content": text},
+                "parent_tool_use_id": None,
+                "session_id": "",
+            }
+
+        return stream()
+
+    await _run_query_with_only_can_use_tool(make_stream)
+
+
+@pytest.mark.e2e
+@pytest.mark.anyio
+async def test_query_string_prompt_with_can_use_tool():
+    """query() accepts a plain string prompt together with can_use_tool."""
+    await _run_query_with_only_can_use_tool(lambda text: text)

src/claude_agent_sdk/_internal/client.py

@@ -3,15 +3,14 @@
 import json
 import os
 from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
-from dataclasses import asdict, replace
+from dataclasses import asdict
 from typing import Any
 
 from ..types import (
     ClaudeAgentOptions,
-    HookEvent,
-    HookMatcher,
     Message,
-    _warn_if_can_use_tool_shadowed,
+    _configure_can_use_tool,
+    _hooks_to_internal_format,
 )
 from .message_parser import parse_message
 from .query import Query
@@ -32,24 +31,6 @@ class InternalClient:
     def __init__(self) -> None:
         """Initialize the internal client."""
 
-    def _convert_hooks_to_internal_format(
-        self, hooks: dict[HookEvent, list[HookMatcher]]
-    ) -> dict[str, list[dict[str, Any]]]:
-        """Convert HookMatcher format to internal Query format."""
-        internal_hooks: dict[str, list[dict[str, Any]]] = {}
-        for event, matchers in hooks.items():
-            internal_hooks[event] = []
-            for matcher in matchers:
-                # Convert HookMatcher to internal dict format
-                internal_matcher: dict[str, Any] = {
-                    "matcher": matcher.matcher if hasattr(matcher, "matcher") else None,
-                    "hooks": matcher.hooks if hasattr(matcher, "hooks") else [],
-                }
-                if hasattr(matcher, "timeout") and matcher.timeout is not None:
-                    internal_matcher["timeout"] = matcher.timeout
-                internal_hooks[event].append(internal_matcher)
-        return internal_hooks
-
     async def process_query(
         self,
         prompt: str | AsyncIterable[dict[str, Any]],
@@ -97,27 +78,7 @@ async def _process_query_inner(
         materialized: MaterializedResume | None,
     ) -> AsyncGenerator[Message, None]:
         # Validate and configure permission settings (matching TypeScript SDK logic)
-        configured_options = options
-        if options.can_use_tool:
-            # canUseTool callback requires streaming mode (AsyncIterable prompt)
-            if isinstance(prompt, str):
-                raise ValueError(
-                    "can_use_tool callback requires streaming mode. "
-                    "Please provide prompt as an AsyncIterable instead of a string."
-                )
-
-            # canUseTool and permission_prompt_tool_name are mutually exclusive
-            if options.permission_prompt_tool_name:
-                raise ValueError(
-                    "can_use_tool callback cannot be used with permission_prompt_tool_name. "
-                    "Please use one or the other."
-                )
-
-            # Advisory: warn if other options shadow the callback
-            _warn_if_can_use_tool_shadowed(options)
-
-            # Automatically set permission_prompt_tool_name to "stdio" for control protocol
-            configured_options = replace(options, permission_prompt_tool_name="stdio")
+        configured_options = _configure_can_use_tool(options)
 
         if materialized is not None:
             configured_options = apply_materialized_options(
@@ -175,7 +136,7 @@ async def _process_query_inner(
             transport=chosen_transport,
             is_streaming_mode=True,  # Always streaming internally
             can_use_tool=configured_options.can_use_tool,
-            hooks=self._convert_hooks_to_internal_format(configured_options.hooks)
+            hooks=_hooks_to_internal_format(configured_options.hooks)
             if configured_options.hooks
             else None,
             sdk_mcp_servers=sdk_mcp_servers,

src/claude_agent_sdk/_internal/query.py

@@ -915,25 +915,45 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None:
             if status in TERMINAL_TASK_STATUSES:
                 self._inflight_tasks.discard(task_id)
 
+    def _has_bidirectional_needs(self) -> bool:
+        """Whether the CLI may still send control requests that need a reply.
+
+        SDK MCP servers, hooks, and the ``can_use_tool`` permission callback
+        are all served over the control protocol: the CLI writes a
+        ``control_request`` to stdout and blocks until the SDK writes the
+        matching ``control_response`` to stdin. Closing stdin while any of
+        these are configured makes every later request fail CLI-side with
+        "Stream closed". Mirrors the TypeScript SDK's ``hasBidirectionalNeeds``.
+        """
+        return bool(self.sdk_mcp_servers or self.hooks or self.can_use_tool)
+
     async def wait_for_result_and_end_input(self) -> None:
-        """Wait for a run-ending result (if needed) then close stdin.
-
-        If SDK MCP servers or hooks require bidirectional communication,
-        keeps stdin open until a result arrives with no tasks in flight. A
-        result frame ends one turn, not necessarily the run: background tasks
-        keep running past it and still need stdin for hook/SDK-MCP control
-        responses (see #1088). The control protocol requires stdin to remain
-        open for the entire conversation, so no timeout is applied. The event
-        is guaranteed to fire: either when a result message arrives with no
-        in-flight tasks (every task completion wakes the parent for a
-        follow-up turn, which ends in such a result), or in _read_messages'
-        finally block if the process exits early.
+        """Wait for the closing result (if needed) then close stdin.
+
+        If SDK MCP servers, hooks, or a ``can_use_tool`` callback require
+        bidirectional communication, keeps stdin open until the first result
+        frame that arrives with no tasks in flight. A result frame ends one
+        turn, not necessarily the run: background tasks keep running past it
+        and still need stdin for control responses (see #1088). The control
+        protocol requires stdin to remain open for the entire conversation, so
+        no timeout is applied. The event is guaranteed to fire: either when a
+        result message arrives with no in-flight tasks (every task completion
+        wakes the parent for a follow-up turn, which ends in such a result),
+        or in _read_messages' finally block if the process exits early.
+
+        Known limitation: the event is one-shot and is not aware of prompt
+        messages still queued CLI-side, so an ``AsyncIterable`` prompt that
+        yields several user messages (several turns) releases the hold at the
+        first turn boundary with no tracked tasks; control requests from later
+        turns can then find stdin closed. Single-message and string prompts —
+        the common one-shot shapes — are fully covered.
         """
-        if self.sdk_mcp_servers or self.hooks:
+        if self._has_bidirectional_needs():
             logger.debug(
                 "Waiting for a run-ending result before closing stdin "
                 f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, "
-                f"has_hooks={bool(self.hooks)})"
+                f"has_hooks={bool(self.hooks)}, "
+                f"has_can_use_tool={self.can_use_tool is not None})"
             )
             await self._first_result_event.wait()
 
@@ -942,18 +962,33 @@ async def wait_for_result_and_end_input(self) -> None:
     async def stream_input(self, stream: AsyncIterable[dict[str, Any]]) -> None:
         """Stream input messages to transport.
 
-        If SDK MCP servers or hooks are present, waits for the first result
-        before closing stdin to allow bidirectional control protocol communication.
+        If SDK MCP servers, hooks, or a ``can_use_tool`` callback are present,
+        waits for a run-ending result before closing stdin to allow
+        bidirectional control protocol communication.
         """
+        written = 0
         try:
             async for message in stream:
                 if self._closed:
                     break
                 await self.transport.write(json.dumps(message) + "\n")
-
-            await self.wait_for_result_and_end_input()
+                written += 1
+        except Exception as e:
+            # A user-supplied prompt iterable (or the write) failed. Don't
+            # leave stdin open — the CLI would wait for input forever and the
+            # consumer's `async for` would never finish — fall through and
+            # close it like a normal end of input.
+            logger.error("Prompt stream failed; closing stdin: %s", e)
+        try:
+            if written:
+                await self.wait_for_result_and_end_input()
+            else:
+                # Nothing was sent, so no result will arrive to release the
+                # hold; close immediately (mirrors the TypeScript SDK's
+                # messageCount guard).
+                await self.transport.end_input()
         except Exception as e:
-            logger.debug(f"Error streaming input: {e}")
+            logger.debug(f"Error closing input stream: {e}")
 
     async def receive_messages(self) -> AsyncIterator[dict[str, Any]]:
         """Receive SDK messages (not control messages)."""

src/claude_agent_sdk/client.py

@@ -3,7 +3,7 @@
 import json
 import os
 from collections.abc import AsyncIterable, AsyncIterator
-from dataclasses import asdict, replace
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any
 
 from . import Transport
@@ -14,13 +14,12 @@
 from .types import (
     ClaudeAgentOptions,
     ContextUsageResponse,
-    HookEvent,
-    HookMatcher,
     McpStatusResponse,
     Message,
     PermissionMode,
     ResultMessage,
-    _warn_if_can_use_tool_shadowed,
+    _configure_can_use_tool,
+    _hooks_to_internal_format,
 )
 
 
@@ -79,24 +78,6 @@ def __init__(
         self._query: Any | None = None
         self._materialized: MaterializedResume | None = None
 
-    def _convert_hooks_to_internal_format(
-        self, hooks: dict[HookEvent, list[HookMatcher]]
-    ) -> dict[str, list[dict[str, Any]]]:
-        """Convert HookMatcher format to internal Query format."""
-        internal_hooks: dict[str, list[dict[str, Any]]] = {}
-        for event, matchers in hooks.items():
-            internal_hooks[event] = []
-            for matcher in matchers:
-                # Convert HookMatcher to internal dict format
-                internal_matcher: dict[str, Any] = {
-                    "matcher": matcher.matcher if hasattr(matcher, "matcher") else None,
-                    "hooks": matcher.hooks if hasattr(matcher, "hooks") else [],
-                }
-                if hasattr(matcher, "timeout") and matcher.timeout is not None:
-                    internal_matcher["timeout"] = matcher.timeout
-                internal_hooks[event].append(internal_matcher)
-        return internal_hooks
-
     async def connect(
         self, prompt: str | AsyncIterable[dict[str, Any]] | None = None
     ) -> None:
@@ -158,28 +139,7 @@ async def _connect_inner(
         from ._internal.transport.subprocess_cli import SubprocessCLITransport
 
         # Validate and configure permission settings (matching TypeScript SDK logic)
-        if self.options.can_use_tool:
-            # canUseTool callback requires streaming mode (AsyncIterable prompt)
-            if isinstance(prompt, str):
-                raise ValueError(
-                    "can_use_tool callback requires streaming mode. "
-                    "Please provide prompt as an AsyncIterable instead of a string."
-                )
-
-            # canUseTool and permission_prompt_tool_name are mutually exclusive
-            if self.options.permission_prompt_tool_name:
-                raise ValueError(
-                    "can_use_tool callback cannot be used with permission_prompt_tool_name. "
-                    "Please use one or the other."
-                )
-
-            # Advisory: warn if other options shadow the callback
-            _warn_if_can_use_tool_shadowed(self.options)
-
-            # Automatically set permission_prompt_tool_name to "stdio" for control protocol
-            options = replace(self.options, permission_prompt_tool_name="stdio")
-        else:
-            options = self.options
+        options = _configure_can_use_tool(self.options)
 
         if self._materialized is not None:
             options = apply_materialized_options(options, self._materialized)
@@ -230,7 +190,7 @@ async def _connect_inner(
             transport=self._transport,
             is_streaming_mode=True,  # ClaudeSDKClient always uses streaming mode
             can_use_tool=self.options.can_use_tool,
-            hooks=self._convert_hooks_to_internal_format(self.options.hooks)
+            hooks=_hooks_to_internal_format(self.options.hooks)
             if self.options.hooks
             else None,
             sdk_mcp_servers=sdk_mcp_servers,