Add agent_id/agent_type to tool-lifecycle hook inputs

src/claude_agent_sdk/types.py

@@ -182,7 +182,32 @@ class BaseHookInput(TypedDict):
     permission_mode: NotRequired[str]
 
 
-class PreToolUseHookInput(BaseHookInput):
+# agent_id/agent_type are present on BaseHookInput in the CLI's schema but are
+# declared per-hook here because SubagentStartHookInput/SubagentStopHookInput
+# need them as *required*, and PEP 655 forbids narrowing NotRequired->Required
+# in a TypedDict subclass. The four tool-lifecycle types below are the only
+# ones the CLI actually populates (the other BaseHookInput consumers don't
+# have a toolUseContext in scope at their build site).
+class _SubagentContextMixin(TypedDict, total=False):
+    """Optional sub-agent attribution fields for tool-lifecycle hooks.
+
+    agent_id: Sub-agent identifier. Present only when the hook fires from
+    inside a Task-spawned sub-agent; absent on the main thread. Matches the
+    agent_id emitted by that sub-agent's SubagentStart/SubagentStop hooks.
+    When multiple sub-agents run in parallel their tool-lifecycle hooks
+    interleave over the same control channel — this is the only reliable
+    way to attribute each one to the correct sub-agent.
+
+    agent_type: Agent type name (e.g. "general-purpose", "code-reviewer").
+    Present inside a sub-agent (alongside agent_id), or on the main thread
+    of a session started with --agent (without agent_id).
+    """
+
+    agent_id: str
+    agent_type: str
+
+
+class PreToolUseHookInput(BaseHookInput, _SubagentContextMixin):
     """Input data for PreToolUse hook events."""
 
     hook_event_name: Literal["PreToolUse"]
@@ -191,7 +216,7 @@ class PreToolUseHookInput(BaseHookInput):
     tool_use_id: str
 
 
-class PostToolUseHookInput(BaseHookInput):
+class PostToolUseHookInput(BaseHookInput, _SubagentContextMixin):
     """Input data for PostToolUse hook events."""
 
     hook_event_name: Literal["PostToolUse"]
@@ -201,7 +226,7 @@ class PostToolUseHookInput(BaseHookInput):
     tool_use_id: str
 
 
-class PostToolUseFailureHookInput(BaseHookInput):
+class PostToolUseFailureHookInput(BaseHookInput, _SubagentContextMixin):
     """Input data for PostToolUseFailure hook events."""
 
     hook_event_name: Literal["PostToolUseFailure"]
@@ -261,7 +286,7 @@ class SubagentStartHookInput(BaseHookInput):
     agent_type: str
 
 
-class PermissionRequestHookInput(BaseHookInput):
+class PermissionRequestHookInput(BaseHookInput, _SubagentContextMixin):
     """Input data for PermissionRequest hook events."""
 
     hook_event_name: Literal["PermissionRequest"]

tests/test_types.py

@@ -203,6 +203,62 @@ def test_subagent_start_hook_input(self):
         assert hook_input["agent_id"] == "agent-42"
         assert hook_input["agent_type"] == "researcher"
 
+    def test_pre_tool_use_hook_input_with_agent_id(self):
+        """PreToolUseHookInput accepts optional agent_id/agent_type.
+
+        When a tool is called from inside a Task sub-agent, the CLI includes
+        the calling agent's id so consumers can correlate the tool call to
+        the correct sub-agent — parallel sub-agents interleave their hook
+        callbacks over the same control channel and are otherwise
+        indistinguishable.
+        """
+        from claude_agent_sdk.types import PreToolUseHookInput
+
+        # Tool called from inside a sub-agent: agent_id present,
+        # same value SubagentStart emits.
+        hook_input: PreToolUseHookInput = {
+            "session_id": "sess-1",
+            "transcript_path": "/tmp/transcript",
+            "cwd": "/home/user",
+            "hook_event_name": "PreToolUse",
+            "tool_name": "Bash",
+            "tool_input": {"command": "echo hello"},
+            "tool_use_id": "toolu_abc123",
+            "agent_id": "agent-42",
+            "agent_type": "researcher",
+        }
+        assert hook_input.get("agent_id") == "agent-42"
+        assert hook_input.get("agent_type") == "researcher"
+
+        # Tool called on the main thread: agent_id absent. Still type-valid.
+        hook_input_main: PreToolUseHookInput = {
+            "session_id": "sess-1",
+            "transcript_path": "/tmp/transcript",
+            "cwd": "/home/user",
+            "hook_event_name": "PreToolUse",
+            "tool_name": "Bash",
+            "tool_input": {"command": "echo hello"},
+            "tool_use_id": "toolu_def456",
+        }
+        assert hook_input_main.get("agent_id") is None
+
+    def test_post_tool_use_hook_input_with_agent_id(self):
+        """PostToolUseHookInput accepts optional agent_id."""
+        from claude_agent_sdk.types import PostToolUseHookInput
+
+        hook_input: PostToolUseHookInput = {
+            "session_id": "sess-1",
+            "transcript_path": "/tmp/transcript",
+            "cwd": "/home/user",
+            "hook_event_name": "PostToolUse",
+            "tool_name": "Bash",
+            "tool_input": {"command": "echo hello"},
+            "tool_response": {"content": [{"type": "text", "text": "hello"}]},
+            "tool_use_id": "toolu_abc123",
+            "agent_id": "agent-42",
+        }
+        assert hook_input.get("agent_id") == "agent-42"
+
     def test_permission_request_hook_input(self):
         """Test PermissionRequestHookInput construction."""
         hook_input: PermissionRequestHookInput = {