Skip to content

Commit 388f2f3

Browse files
Add permission checks for SDK-registered custom tools (#555)
* Add permission checks for SDK-registered custom tools Add 'custom-tool' to the PermissionRequest kind union in Node.js and Python types. Update all existing custom tool e2e tests across all four languages (Node.js, Python, Go, .NET) to provide an onPermissionRequest handler, and add new e2e tests verifying permission approval and denial flows for custom tools. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR review: remove unused import, add toolName verification to Go and .NET tests - Remove unused PermissionRequestResult import from Node.js test - Add toolName assertion in Go test for cross-SDK parity - Add toolName assertion in .NET test for cross-SDK parity Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Formatting * Fix rebase issue * Go fix --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 9d998fb commit 388f2f3

File tree

9 files changed

+358
-3
lines changed

9 files changed

+358
-3
lines changed

dotnet/test/ToolsTests.cs

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using GitHub.Copilot.SDK.Test.Harness;
66
using Microsoft.Extensions.AI;
77
using System.ComponentModel;
8+
using System.Linq;
89
using System.Text.Json;
910
using System.Text.Json.Serialization;
1011
using Xunit;
@@ -42,6 +43,7 @@ public async Task Invokes_Custom_Tool()
4243
var session = await CreateSessionAsync(new SessionConfig
4344
{
4445
Tools = [AIFunctionFactory.Create(EncryptString, "encrypt_string")],
46+
OnPermissionRequest = PermissionHandler.ApproveAll,
4547
});
4648

4749
await session.SendAsync(new MessageOptions
@@ -66,7 +68,8 @@ public async Task Handles_Tool_Calling_Errors()
6668

6769
var session = await CreateSessionAsync(new SessionConfig
6870
{
69-
Tools = [getUserLocation]
71+
Tools = [getUserLocation],
72+
OnPermissionRequest = PermissionHandler.ApproveAll,
7073
});
7174

7275
await session.SendAsync(new MessageOptions { Prompt = "What is my location? If you can't find out, just say 'unknown'." });
@@ -108,6 +111,7 @@ public async Task Can_Receive_And_Return_Complex_Types()
108111
var session = await CreateSessionAsync(new SessionConfig
109112
{
110113
Tools = [AIFunctionFactory.Create(PerformDbQuery, "db_query", serializerOptions: ToolsTestsJsonContext.Default.Options)],
114+
OnPermissionRequest = PermissionHandler.ApproveAll,
111115
});
112116

113117
await session.SendAsync(new MessageOptions
@@ -154,6 +158,7 @@ public async Task Can_Return_Binary_Result()
154158
var session = await CreateSessionAsync(new SessionConfig
155159
{
156160
Tools = [AIFunctionFactory.Create(GetImage, "get_image")],
161+
OnPermissionRequest = PermissionHandler.ApproveAll,
157162
});
158163

159164
await session.SendAsync(new MessageOptions
@@ -177,4 +182,72 @@ await session.SendAsync(new MessageOptions
177182
SessionLog = "Returned an image",
178183
});
179184
}
185+
186+
[Fact]
187+
public async Task Invokes_Custom_Tool_With_Permission_Handler()
188+
{
189+
var permissionRequests = new List<PermissionRequest>();
190+
191+
var session = await Client.CreateSessionAsync(new SessionConfig
192+
{
193+
Tools = [AIFunctionFactory.Create(EncryptStringForPermission, "encrypt_string")],
194+
OnPermissionRequest = (request, invocation) =>
195+
{
196+
permissionRequests.Add(request);
197+
return Task.FromResult(new PermissionRequestResult { Kind = "approved" });
198+
},
199+
});
200+
201+
await session.SendAsync(new MessageOptions
202+
{
203+
Prompt = "Use encrypt_string to encrypt this string: Hello"
204+
});
205+
206+
var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
207+
Assert.NotNull(assistantMessage);
208+
Assert.Contains("HELLO", assistantMessage!.Data.Content ?? string.Empty);
209+
210+
// Should have received a custom-tool permission request with the correct tool name
211+
var customToolRequest = permissionRequests.FirstOrDefault(r => r.Kind == "custom-tool");
212+
Assert.NotNull(customToolRequest);
213+
Assert.True(customToolRequest!.ExtensionData?.ContainsKey("toolName") ?? false);
214+
var toolName = ((JsonElement)customToolRequest.ExtensionData!["toolName"]).GetString();
215+
Assert.Equal("encrypt_string", toolName);
216+
217+
[Description("Encrypts a string")]
218+
static string EncryptStringForPermission([Description("String to encrypt")] string input)
219+
=> input.ToUpperInvariant();
220+
}
221+
222+
[Fact]
223+
public async Task Denies_Custom_Tool_When_Permission_Denied()
224+
{
225+
var toolHandlerCalled = false;
226+
227+
var session = await Client.CreateSessionAsync(new SessionConfig
228+
{
229+
Tools = [AIFunctionFactory.Create(EncryptStringDenied, "encrypt_string")],
230+
OnPermissionRequest = (request, invocation) =>
231+
{
232+
return Task.FromResult(new PermissionRequestResult { Kind = "denied-interactively-by-user" });
233+
},
234+
});
235+
236+
await session.SendAsync(new MessageOptions
237+
{
238+
Prompt = "Use encrypt_string to encrypt this string: Hello"
239+
});
240+
241+
await TestHelper.GetFinalAssistantMessageAsync(session);
242+
243+
// The tool handler should NOT have been called since permission was denied
244+
Assert.False(toolHandlerCalled);
245+
246+
[Description("Encrypts a string")]
247+
string EncryptStringDenied([Description("String to encrypt")] string input)
248+
{
249+
toolHandlerCalled = true;
250+
return input.ToUpperInvariant();
251+
}
252+
}
180253
}

go/internal/e2e/tools_test.go

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"os"
66
"path/filepath"
77
"strings"
8+
"sync"
89
"testing"
910

1011
copilot "github.com/github/copilot-sdk/go"
@@ -262,4 +263,103 @@ func TestTools(t *testing.T) {
262263
t.Errorf("Expected session ID '%s', got '%s'", session.SessionID, receivedInvocation.SessionID)
263264
}
264265
})
266+
267+
t.Run("invokes custom tool with permission handler", func(t *testing.T) {
268+
ctx.ConfigureForTest(t)
269+
270+
type EncryptParams struct {
271+
Input string `json:"input" jsonschema:"String to encrypt"`
272+
}
273+
274+
var permissionRequests []copilot.PermissionRequest
275+
var mu sync.Mutex
276+
277+
session, err := client.CreateSession(t.Context(), &copilot.SessionConfig{
278+
Tools: []copilot.Tool{
279+
copilot.DefineTool("encrypt_string", "Encrypts a string",
280+
func(params EncryptParams, inv copilot.ToolInvocation) (string, error) {
281+
return strings.ToUpper(params.Input), nil
282+
}),
283+
},
284+
OnPermissionRequest: func(request copilot.PermissionRequest, invocation copilot.PermissionInvocation) (copilot.PermissionRequestResult, error) {
285+
mu.Lock()
286+
permissionRequests = append(permissionRequests, request)
287+
mu.Unlock()
288+
return copilot.PermissionRequestResult{Kind: "approved"}, nil
289+
},
290+
})
291+
if err != nil {
292+
t.Fatalf("Failed to create session: %v", err)
293+
}
294+
295+
_, err = session.Send(t.Context(), copilot.MessageOptions{Prompt: "Use encrypt_string to encrypt this string: Hello"})
296+
if err != nil {
297+
t.Fatalf("Failed to send message: %v", err)
298+
}
299+
300+
answer, err := testharness.GetFinalAssistantMessage(t.Context(), session)
301+
if err != nil {
302+
t.Fatalf("Failed to get assistant message: %v", err)
303+
}
304+
305+
if answer.Data.Content == nil || !strings.Contains(*answer.Data.Content, "HELLO") {
306+
t.Errorf("Expected answer to contain 'HELLO', got %v", answer.Data.Content)
307+
}
308+
309+
// Should have received a custom-tool permission request
310+
mu.Lock()
311+
customToolReqs := 0
312+
for _, req := range permissionRequests {
313+
if req.Kind == "custom-tool" {
314+
customToolReqs++
315+
if toolName, ok := req.Extra["toolName"].(string); !ok || toolName != "encrypt_string" {
316+
t.Errorf("Expected toolName 'encrypt_string', got '%v'", req.Extra["toolName"])
317+
}
318+
}
319+
}
320+
mu.Unlock()
321+
if customToolReqs == 0 {
322+
t.Errorf("Expected at least one custom-tool permission request, got none")
323+
}
324+
})
325+
326+
t.Run("denies custom tool when permission denied", func(t *testing.T) {
327+
ctx.ConfigureForTest(t)
328+
329+
type EncryptParams struct {
330+
Input string `json:"input" jsonschema:"String to encrypt"`
331+
}
332+
333+
toolHandlerCalled := false
334+
335+
session, err := client.CreateSession(t.Context(), &copilot.SessionConfig{
336+
Tools: []copilot.Tool{
337+
copilot.DefineTool("encrypt_string", "Encrypts a string",
338+
func(params EncryptParams, inv copilot.ToolInvocation) (string, error) {
339+
toolHandlerCalled = true
340+
return strings.ToUpper(params.Input), nil
341+
}),
342+
},
343+
OnPermissionRequest: func(request copilot.PermissionRequest, invocation copilot.PermissionInvocation) (copilot.PermissionRequestResult, error) {
344+
return copilot.PermissionRequestResult{Kind: "denied-interactively-by-user"}, nil
345+
},
346+
})
347+
if err != nil {
348+
t.Fatalf("Failed to create session: %v", err)
349+
}
350+
351+
_, err = session.Send(t.Context(), copilot.MessageOptions{Prompt: "Use encrypt_string to encrypt this string: Hello"})
352+
if err != nil {
353+
t.Fatalf("Failed to send message: %v", err)
354+
}
355+
356+
_, err = testharness.GetFinalAssistantMessage(t.Context(), session)
357+
if err != nil {
358+
t.Fatalf("Failed to get assistant message: %v", err)
359+
}
360+
361+
if toolHandlerCalled {
362+
t.Errorf("Tool handler should NOT have been called since permission was denied")
363+
}
364+
})
265365
}

go/types.go

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,32 @@ type PermissionRequest struct {
106106
Extra map[string]any `json:"-"` // Additional fields vary by kind
107107
}
108108

109+
// UnmarshalJSON implements custom JSON unmarshaling for PermissionRequest
110+
// to capture additional fields (varying by kind) into the Extra map.
111+
func (p *PermissionRequest) UnmarshalJSON(data []byte) error {
112+
// Unmarshal known fields via an alias to avoid infinite recursion
113+
type Alias PermissionRequest
114+
var alias Alias
115+
if err := json.Unmarshal(data, &alias); err != nil {
116+
return err
117+
}
118+
*p = PermissionRequest(alias)
119+
120+
// Unmarshal all fields into a generic map
121+
var raw map[string]any
122+
if err := json.Unmarshal(data, &raw); err != nil {
123+
return err
124+
}
125+
126+
// Remove known fields, keep the rest as Extra
127+
delete(raw, "kind")
128+
delete(raw, "toolCallId")
129+
if len(raw) > 0 {
130+
p.Extra = raw
131+
}
132+
return nil
133+
}
134+
109135
// PermissionRequestResult represents the result of a permission request
110136
type PermissionRequestResult struct {
111137
Kind string `json:"kind"`

nodejs/src/types.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,7 @@ export type SystemMessageConfig = SystemMessageAppendConfig | SystemMessageRepla
211211
* Permission request types from the server
212212
*/
213213
export interface PermissionRequest {
214-
kind: "shell" | "write" | "mcp" | "read" | "url";
214+
kind: "shell" | "write" | "mcp" | "read" | "url" | "custom-tool";
215215
toolCallId?: string;
216216
[key: string]: unknown;
217217
}

nodejs/test/e2e/tools.test.ts

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { join } from "path";
77
import { assert, describe, expect, it } from "vitest";
88
import { z } from "zod";
99
import { defineTool, approveAll } from "../../src/index.js";
10+
import type { PermissionRequest } from "../../src/index.js";
1011
import { createSdkTestContext } from "./harness/sdkTestContext";
1112

1213
describe("Custom tools", async () => {
@@ -36,6 +37,7 @@ describe("Custom tools", async () => {
3637
handler: ({ input }) => input.toUpperCase(),
3738
}),
3839
],
40+
onPermissionRequest: approveAll,
3941
});
4042

4143
const assistantMessage = await session.sendAndWait({
@@ -55,6 +57,7 @@ describe("Custom tools", async () => {
5557
},
5658
}),
5759
],
60+
onPermissionRequest: approveAll,
5861
});
5962

6063
const answer = await session.sendAndWait({
@@ -111,6 +114,7 @@ describe("Custom tools", async () => {
111114
},
112115
}),
113116
],
117+
onPermissionRequest: approveAll,
114118
});
115119

116120
const assistantMessage = await session.sendAndWait({
@@ -127,4 +131,63 @@ describe("Custom tools", async () => {
127131
expect(responseContent.replace(/,/g, "")).toContain("135460");
128132
expect(responseContent.replace(/,/g, "")).toContain("204356");
129133
});
134+
135+
it("invokes custom tool with permission handler", async () => {
136+
const permissionRequests: PermissionRequest[] = [];
137+
138+
const session = await client.createSession({
139+
tools: [
140+
defineTool("encrypt_string", {
141+
description: "Encrypts a string",
142+
parameters: z.object({
143+
input: z.string().describe("String to encrypt"),
144+
}),
145+
handler: ({ input }) => input.toUpperCase(),
146+
}),
147+
],
148+
onPermissionRequest: (request) => {
149+
permissionRequests.push(request);
150+
return { kind: "approved" };
151+
},
152+
});
153+
154+
const assistantMessage = await session.sendAndWait({
155+
prompt: "Use encrypt_string to encrypt this string: Hello",
156+
});
157+
expect(assistantMessage?.data.content).toContain("HELLO");
158+
159+
// Should have received a custom-tool permission request
160+
const customToolRequests = permissionRequests.filter((req) => req.kind === "custom-tool");
161+
expect(customToolRequests.length).toBeGreaterThan(0);
162+
expect(customToolRequests[0].toolName).toBe("encrypt_string");
163+
});
164+
165+
it("denies custom tool when permission denied", async () => {
166+
let toolHandlerCalled = false;
167+
168+
const session = await client.createSession({
169+
tools: [
170+
defineTool("encrypt_string", {
171+
description: "Encrypts a string",
172+
parameters: z.object({
173+
input: z.string().describe("String to encrypt"),
174+
}),
175+
handler: ({ input }) => {
176+
toolHandlerCalled = true;
177+
return input.toUpperCase();
178+
},
179+
}),
180+
],
181+
onPermissionRequest: () => {
182+
return { kind: "denied-interactively-by-user" };
183+
},
184+
});
185+
186+
await session.sendAndWait({
187+
prompt: "Use encrypt_string to encrypt this string: Hello",
188+
});
189+
190+
// The tool handler should NOT have been called since permission was denied
191+
expect(toolHandlerCalled).toBe(false);
192+
});
130193
});

python/copilot/types.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ class SystemMessageReplaceConfig(TypedDict):
169169
class PermissionRequest(TypedDict, total=False):
170170
"""Permission request from the server"""
171171

172-
kind: Literal["shell", "write", "mcp", "read", "url"]
172+
kind: Literal["shell", "write", "mcp", "read", "url", "custom-tool"]
173173
toolCallId: str
174174
# Additional fields vary by kind
175175

0 commit comments

Comments
 (0)