diff --git a/src/main/java/com/github/copilot/sdk/CopilotSession.java b/src/main/java/com/github/copilot/sdk/CopilotSession.java index b54b2cbf3..57757c4b4 100644 --- a/src/main/java/com/github/copilot/sdk/CopilotSession.java +++ b/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. *
diff --git a/src/main/java/com/github/copilot/sdk/RpcHandlerDispatcher.java b/src/main/java/com/github/copilot/sdk/RpcHandlerDispatcher.java index d085f7fce..5cb0aea8c 100644 --- a/src/main/java/com/github/copilot/sdk/RpcHandlerDispatcher.java +++ b/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,48 @@ private void handleSystemMessageTransform(JsonRpcClient rpc, String requestId, J }); } + /** + * Finds the first registered session that has a hooks handler registered. + *
+ * Used as a fallback when the session ID in an incoming {@code hooks.invoke} + * request is not in the registry (e.g. sub-agent sessions created internally by + * the CLI). When multiple sessions have hooks registered, the selection is + * non-deterministic because the underlying + * {@link java.util.concurrent.ConcurrentHashMap} does not guarantee iteration + * order. + * + * @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. + *
+ * Used as a fallback when the session ID in an incoming + * {@code permission.request} request is not in the registry (e.g. sub-agent + * sessions created internally by the CLI). When multiple sessions have + * permission handlers registered, the selection is non-deterministic because + * the underlying {@link java.util.concurrent.ConcurrentHashMap} does not + * guarantee iteration order. + * + * @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) { diff --git a/src/test/java/com/github/copilot/sdk/RpcHandlerDispatcherTest.java b/src/test/java/com/github/copilot/sdk/RpcHandlerDispatcherTest.java index 315a38b90..505aee670 100644 --- a/src/test/java/com/github/copilot/sdk/RpcHandlerDispatcherTest.java +++ b/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