Fix onPreToolUse hook for sub-agent sessions via fallback dispatch

Agent-Logs-Url: https://github.com/github/copilot-sdk-java/sessions/4ea5c73c-1940-4862-8afb-019a8de0bca2

Co-authored-by: edburns <75821+edburns@users.noreply.github.com>

Author: Copilot

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

@@ -1219,6 +1219,25 @@ void registerPermissionHandler(PermissionHandler handler) {
         permissionHandler.set(handler);
     }
 
+    /**
+     * Returns whether this session has a permission handler registered.
+     *
+     * @return {@code true} if a permission handler is set
+     */
+    boolean hasPermissionHandler() {
+        return permissionHandler.get() != null;
+    }
+
+    /**
+     * Returns whether this session has any hooks registered.
+     *
+     * @return {@code true} if a hooks handler with at least one hook is set
+     */
+    boolean hasHooksHandler() {
+        SessionHooks hooks = hooksHandler.get();
+        return hooks != null && hooks.hasHooks();
+    }
+
     /**
      * Handles a permission request from the Copilot CLI.
      * <p>

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

@@ -190,6 +190,12 @@ private void handlePermissionRequest(JsonRpcClient rpc, String requestId, JsonNo
                 JsonNode permissionRequest = params.get("permissionRequest");
 
                 CopilotSession session = sessions.get(sessionId);
+                if (session == null) {
+                    // Fall back to any registered session that has a permission handler.
+                    // This handles sub-agent sessions created internally by the CLI whose IDs
+                    // are not in the SDK's session registry.
+                    session = findSessionWithPermissionHandler();
+                }
                 if (session == null) {
                     var result = new PermissionRequestResult()
                             .setKind(PermissionRequestResultKind.DENIED_COULD_NOT_REQUEST_FROM_USER);
@@ -292,7 +298,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);
+                    // Fall back to any registered session that has hooks registered.
+                    // This handles sub-agent sessions created internally by the CLI whose IDs
+                    // are not in the SDK's session registry.
+                    session = findSessionWithHooks();
+                }
+                if (session == null) {
+                    // No session with hooks — return null output (no-op) rather than an error.
+                    rpc.sendResponse(Long.parseLong(requestId), Collections.singletonMap("output", null));
                     return;
                 }
 
@@ -366,6 +379,34 @@ private void handleSystemMessageTransform(JsonRpcClient rpc, String requestId, J
         });
     }
 
+    /**
+     * Finds the first registered session that has a hooks handler registered.
+     *
+     * @return a session with hooks, or {@code null} if none is found
+     */
+    private CopilotSession findSessionWithHooks() {
+        for (CopilotSession s : sessions.values()) {
+            if (s.hasHooksHandler()) {
+                return s;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Finds the first registered session that has a permission handler registered.
+     *
+     * @return a session with a permission handler, or {@code null} if none is found
+     */
+    private CopilotSession findSessionWithPermissionHandler() {
+        for (CopilotSession s : sessions.values()) {
+            if (s.hasPermissionHandler()) {
+                return s;
+            }
+        }
+        return null;
+    }
+
     private void runAsync(Runnable task) {
         try {
             if (executor != null) {

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 registered → denied
         ObjectNode params = MAPPER.createObjectNode();
         params.put("sessionId", "nonexistent");
         params.putObject("permissionRequest");
@@ -307,6 +308,25 @@ void permissionRequestWithUnknownSession() throws Exception {
         assertEquals("denied-no-approval-rule-and-could-not-request-from-user", result.get("kind").asText());
     }
 
+    @Test
+    void permissionRequestWithSubAgentSessionFallsBackToParentWithHandler() throws Exception {
+        // Register a parent session with a permission handler, invoke with a sub-agent
+        // session ID
+        CopilotSession parentSession = createSession("parent-session");
+        parentSession.registerPermissionHandler((request, invocation) -> CompletableFuture
+                .completedFuture(new PermissionRequestResult().setKind("allow")));
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "sub-agent-session-id"); // Not in registry
+        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 +473,8 @@ void userInputRequestHandlerFails() throws Exception {
     // ===== hooks.invoke tests =====
 
     @Test
-    void hooksInvokeWithUnknownSession() throws Exception {
+    void hooksInvokeWithUnknownSessionAndNoFallback() throws Exception {
+        // No sessions registered at all → returns null output (no-op)
         ObjectNode params = MAPPER.createObjectNode();
         params.put("sessionId", "nonexistent");
         params.put("hookType", "preToolUse");
@@ -462,8 +483,52 @@ void hooksInvokeWithUnknownSession() throws Exception {
         invokeHandler("hooks.invoke", "30", params);
 
         JsonNode response = readResponse();
-        assertNotNull(response.get("error"));
-        assertEquals(-32602, response.get("error").get("code").asInt());
+        // No sessions with hooks → null output, not an error
+        assertNull(response.get("error"), "Should not return an error when no fallback session exists");
+        JsonNode output = response.get("result").get("output");
+        assertTrue(output == null || output.isNull(), "Output should be null when no session with hooks is found");
+    }
+
+    @Test
+    void hooksInvokeWithSubAgentSessionFallsBackToParentWithHooks() throws Exception {
+        // Register a parent session with hooks, but invoke with a sub-agent session ID
+        CopilotSession parentSession = createSession("parent-session");
+        parentSession.registerHooks(new SessionHooks().setOnPreToolUse(
+                (input, invocation) -> CompletableFuture.completedFuture(PreToolUseHookOutput.allow())));
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "sub-agent-session-id"); // Not in registry
+        params.put("hookType", "preToolUse");
+        ObjectNode input = params.putObject("input");
+        input.put("toolName", "glob");
+        input.put("toolCallId", "tc-sub-1");
+
+        invokeHandler("hooks.invoke", "35", params);
+
+        JsonNode response = readResponse();
+        assertNull(response.get("error"), "Should not return an error when a fallback session with hooks exists");
+        JsonNode output = response.get("result").get("output");
+        assertNotNull(output);
+        assertEquals("allow", output.get("permissionDecision").asText());
+    }
+
+    @Test
+    void hooksInvokeWithSubAgentSessionAndNoSessionHasHooks() throws Exception {
+        // Register a parent session WITHOUT hooks, invoke with unknown sub-agent
+        // session ID
+        createSession("parent-session-no-hooks");
+
+        ObjectNode params = MAPPER.createObjectNode();
+        params.put("sessionId", "sub-agent-session-id"); // Not in registry
+        params.put("hookType", "preToolUse");
+        params.putObject("input");
+
+        invokeHandler("hooks.invoke", "36", params);
+
+        JsonNode response = readResponse();
+        assertNull(response.get("error"), "Should not return an error when no fallback session has hooks");
+        JsonNode output = response.get("result").get("output");
+        assertTrue(output == null || output.isNull(), "Output should be null when no session with hooks is found");
     }
 
     @Test