Fix onPreToolUse hook not firing for sub-agent tool calls

When the CLI creates sub-agent sessions internally, their session IDs
are not in the SDK's session registry. The hooks.invoke and
permission.request handlers now fall back to any registered session
that has the appropriate handler, enabling hooks and permission checks
to fire for sub-agent tool calls.

Agent-Logs-Url: https://github.com/github/copilot-sdk-java/sessions/868c6fed-c57d-4d9f-806c-eca509096672

Co-authored-by: brunoborges <129743+brunoborges@users.noreply.github.com>

Author: Copilot

src/main/java/com/github/copilot/sdk/CopilotSession.java

@@ -1298,6 +1298,31 @@ void registerHooks(SessionHooks hooks) {
         hooksHandler.set(hooks);
     }
 
+    /**
+     * Returns whether this session has hooks registered.
+     * <p>
+     * Used internally to find a session that can handle hooks invocations for
+     * sub-agent sessions whose IDs are not directly tracked by the SDK.
+     *
+     * @return {@code true} if hooks are registered and at least one handler is set
+     */
+    boolean hasHooks() {
+        SessionHooks hooks = hooksHandler.get();
+        return hooks != null && hooks.hasHooks();
+    }
+
+    /**
+     * Returns whether this session has a permission handler registered.
+     * <p>
+     * Used internally to find a session that can handle permission requests for
+     * sub-agent sessions whose IDs are not directly tracked by the SDK.
+     *
+     * @return {@code true} if a permission handler is registered
+     */
+    boolean hasPermissionHandler() {
+        return permissionHandler.get() != null;
+    }
+
     /**
      * Registers transform callbacks for system message sections.
      * <p>

src/main/java/com/github/copilot/sdk/RpcHandlerDispatcher.java

@@ -191,6 +191,12 @@ private void handlePermissionRequest(JsonRpcClient rpc, String requestId, JsonNo
                 JsonNode permissionRequest = params.get("permissionRequest");
 
                 CopilotSession session = sessions.get(sessionId);
+                if (session == null) {
+                    // The CLI may send permission requests for sub-agent sessions
+                    // whose IDs are not tracked by the SDK. Fall back to any
+                    // registered session that has a permission handler.
+                    session = findSessionWithPermissionHandler();
+                }
                 if (session == null) {
                     var result = new PermissionRequestResult()
                             .setKind(PermissionRequestResultKind.DENIED_COULD_NOT_REQUEST_FROM_USER);
@@ -293,7 +299,14 @@ private void handleHooksInvoke(JsonRpcClient rpc, String requestId, JsonNode par
 
                 CopilotSession session = sessions.get(sessionId);
                 if (session == null) {
-                    rpc.sendErrorResponse(Long.parseLong(requestId), -32602, "Unknown session " + sessionId);
+                    // The CLI may send hooks.invoke for sub-agent sessions whose IDs
+                    // are not tracked by the SDK. Fall back to any registered session
+                    // that has hooks so that application-level hooks (e.g. onPreToolUse)
+                    // also fire for sub-agent tool calls.
+                    session = findSessionWithHooks();
+                }
+                if (session == null) {
+                    rpc.sendResponse(Long.parseLong(requestId), Collections.singletonMap("output", null));
                     return;
                 }
 
@@ -318,6 +331,41 @@ private void handleHooksInvoke(JsonRpcClient rpc, String requestId, JsonNode par
         });
     }
 
+    /**
+     * Finds a session that has hooks registered.
+     * <p>
+     * When the CLI creates sub-agent sessions internally, their session IDs are not
+     * in the SDK's session map. This method searches all registered sessions to
+     * find one with hooks, enabling hook handlers to fire for sub-agent tool calls.
+     *
+     * @return a session with hooks, or {@code null} if none found
+     */
+    private CopilotSession findSessionWithHooks() {
+        for (CopilotSession s : sessions.values()) {
+            if (s.hasHooks()) {
+                return s;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Finds a session that has a permission handler registered.
+     * <p>
+     * Similar to {@link #findSessionWithHooks()}, this enables permission handlers
+     * to fire for sub-agent sessions whose IDs are not tracked by the SDK.
+     *
+     * @return a session with a permission handler, or {@code null} if none found
+     */
+    private CopilotSession findSessionWithPermissionHandler() {
+        for (CopilotSession s : sessions.values()) {
+            if (s.hasPermissionHandler()) {
+                return s;
+            }
+        }
+        return null;
+    }
+
     /**
      * Functional interface for dispatching lifecycle events.
      */

src/test/java/com/github/copilot/sdk/RpcHandlerDispatcherTest.java

@@ -295,7 +295,8 @@ void toolCallHandlerFails() throws Exception {
     // ===== permission.request tests =====
 
     @Test
-    void permissionRequestWithUnknownSession() throws Exception {
+    void permissionRequestWithUnknownSessionAndNoFallback() throws Exception {
+        // No sessions at all — returns denied
         ObjectNode params = MAPPER.createObjectNode();
         params.put("sessionId", "nonexistent");
         params.putObject("permissionRequest");
@@ -307,6 +308,24 @@ void permissionRequestWithUnknownSession() throws Exception {
         assertEquals("denied-no-approval-rule-and-could-not-request-from-user", result.get("kind").asText());
     }
 
+    @Test
+    void permissionRequestFallsBackToSessionWithHandlerForSubAgent() throws Exception {
+        // Parent session has permission handler; sub-agent session ID not in map.
+        CopilotSession parent = createSession("parent-session");
+        parent.registerPermissionHandler((request, invocation) -> CompletableFuture
+                .completedFuture(new PermissionRequestResult().setKind("allow")));
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "subagent-session-id");
+        params.putObject("permissionRequest");
+
+        invokeHandler("permission.request", "15", params);
+
+        JsonNode response = readResponse();
+        JsonNode result = response.get("result").get("result");
+        assertEquals("allow", result.get("kind").asText());
+    }
+
     @Test
     void permissionRequestWithHandler() throws Exception {
         CopilotSession session = createSession("s1");
@@ -453,7 +472,8 @@ void userInputRequestHandlerFails() throws Exception {
     // ===== hooks.invoke tests =====
 
     @Test
-    void hooksInvokeWithUnknownSession() throws Exception {
+    void hooksInvokeWithUnknownSessionAndNoFallback() throws Exception {
+        // No sessions at all — returns null output (no-op)
         ObjectNode params = MAPPER.createObjectNode();
         params.put("sessionId", "nonexistent");
         params.put("hookType", "preToolUse");
@@ -462,8 +482,49 @@ void hooksInvokeWithUnknownSession() throws Exception {
         invokeHandler("hooks.invoke", "30", params);
 
         JsonNode response = readResponse();
-        assertNotNull(response.get("error"));
-        assertEquals(-32602, response.get("error").get("code").asInt());
+        JsonNode output = response.get("result").get("output");
+        assertTrue(output == null || output.isNull(),
+                "Output should be null when no sessions with hooks are registered");
+    }
+
+    @Test
+    void hooksInvokeFallsBackToSessionWithHooksForSubAgent() throws Exception {
+        // Parent session has hooks; sub-agent session ID is not in the map.
+        // The dispatcher should fall back to the parent session's hooks.
+        CopilotSession parent = createSession("parent-session");
+        parent.registerHooks(new SessionHooks().setOnPreToolUse(
+                (input, invocation) -> CompletableFuture.completedFuture(PreToolUseHookOutput.allow())));
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "subagent-session-id");
+        params.put("hookType", "preToolUse");
+        ObjectNode input = params.putObject("input");
+        input.put("toolName", "glob");
+        input.put("toolCallId", "tc-sub");
+
+        invokeHandler("hooks.invoke", "35", params);
+
+        JsonNode response = readResponse();
+        JsonNode output = response.get("result").get("output");
+        assertNotNull(output, "Hooks should fire for sub-agent tool calls via fallback");
+        assertEquals("allow", output.get("permissionDecision").asText());
+    }
+
+    @Test
+    void hooksInvokeFallbackSkipsSessionWithoutHooks() throws Exception {
+        // Session exists but has no hooks — should not be used as fallback
+        createSession("no-hooks-session");
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "subagent-session-id");
+        params.put("hookType", "preToolUse");
+        params.putObject("input");
+
+        invokeHandler("hooks.invoke", "36", params);
+
+        JsonNode response = readResponse();
+        JsonNode output = response.get("result").get("output");
+        assertTrue(output == null || output.isNull(), "Should return null when no session with hooks is found");
     }
 
     @Test