Skip to content

Commit 7a9d57c

Browse files
johnstcnjohnstcn
authored
fix(coderd): actually wire the chat template allowlist into tools (#23626)
Problem: previously, the deployment-wide chat template allowlist was never actually wired in from `chatd.go` - Extracts `parseChatTemplateAllowlist` into shared `coderd/util/xjson.ParseUUIDList` - Adds `Server.chatTemplateAllowlist()` method that reads the allowlist from DB - Passes `AllowedTemplateIDs` callback to `ListTemplates`, `ReadTemplate`, and `CreateWorkspace` tool constructors > 🤖 Created by Coder Agents and reviewed by a human.
1 parent dab4e6f commit 7a9d57c

5 files changed

Lines changed: 262 additions & 23 deletions

File tree

‎coderd/exp_chats.go‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
"github.com/coder/coder/v2/coderd/searchquery"
4545
"github.com/coder/coder/v2/coderd/tracing"
4646
"github.com/coder/coder/v2/coderd/util/ptr"
47+
"github.com/coder/coder/v2/coderd/util/xjson"
4748
"github.com/coder/coder/v2/coderd/workspaceapps"
4849
"github.com/coder/coder/v2/coderd/x/chatd"
4950
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
@@ -2870,14 +2871,18 @@ func (api *API) getChatTemplateAllowlist(rw http.ResponseWriter, r *http.Request
28702871
})
28712872
return
28722873
}
2873-
ids, parseErr := parseChatTemplateAllowlist(raw)
2874+
parsed, parseErr := xjson.ParseUUIDList(raw)
28742875
if parseErr != nil {
28752876
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
28762877
Message: "Stored template allowlist is corrupt.",
28772878
Detail: parseErr.Error(),
28782879
})
28792880
return
28802881
}
2882+
ids := make([]string, len(parsed))
2883+
for i, id := range parsed {
2884+
ids[i] = id.String()
2885+
}
28812886
resp := codersdk.ChatTemplateAllowlist{
28822887
TemplateIDs: ids,
28832888
}
@@ -2983,24 +2988,6 @@ func (api *API) putChatTemplateAllowlist(rw http.ResponseWriter, r *http.Request
29832988
rw.WriteHeader(http.StatusNoContent)
29842989
}
29852990

2986-
// parseChatTemplateAllowlist parses the raw JSON string from the
2987-
// database into a list of template ID strings. Returns an empty
2988-
// slice when the value is empty. Returns an error when the stored
2989-
// JSON is corrupt or otherwise cannot be unmarshalled.
2990-
func parseChatTemplateAllowlist(raw string) ([]string, error) {
2991-
if raw == "" {
2992-
return []string{}, nil
2993-
}
2994-
var ids []string
2995-
if err := json.Unmarshal([]byte(raw), &ids); err != nil {
2996-
return nil, xerrors.Errorf("unmarshal template allowlist: %w", err)
2997-
}
2998-
if ids == nil {
2999-
return []string{}, nil
3000-
}
3001-
return ids, nil
3002-
}
3003-
30042991
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
30052992
//
30062993
//nolint:revive // get-return: revive assumes get* must be a getter, but this is an HTTP handler.

‎coderd/util/xjson/xjson.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package xjson
2+
3+
import (
4+
"encoding/json"
5+
"strings"
6+
7+
"github.com/google/uuid"
8+
"golang.org/x/xerrors"
9+
)
10+
11+
// ParseUUIDList parses a JSON-encoded array of UUID strings
12+
// (e.g. `["uuid1","uuid2"]`) and returns the corresponding
13+
// slice of uuid.UUID values. An empty input (including
14+
// whitespace-only) returns an empty (non-nil) slice.
15+
func ParseUUIDList(raw string) ([]uuid.UUID, error) {
16+
raw = strings.TrimSpace(raw)
17+
if raw == "" {
18+
return []uuid.UUID{}, nil
19+
}
20+
21+
var strs []string
22+
if err := json.Unmarshal([]byte(raw), &strs); err != nil {
23+
return nil, xerrors.Errorf("unmarshal uuid list: %w", err)
24+
}
25+
26+
ids := make([]uuid.UUID, 0, len(strs))
27+
for _, s := range strs {
28+
id, err := uuid.Parse(s)
29+
if err != nil {
30+
return nil, xerrors.Errorf("parse uuid %q: %w", s, err)
31+
}
32+
ids = append(ids, id)
33+
}
34+
return ids, nil
35+
}

‎coderd/util/xjson/xjson_test.go‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package xjson_test
2+
3+
import (
4+
"testing"
5+
6+
"github.com/google/uuid"
7+
"github.com/stretchr/testify/require"
8+
9+
"github.com/coder/coder/v2/coderd/util/xjson"
10+
)
11+
12+
func TestParseUUIDList(t *testing.T) {
13+
t.Parallel()
14+
15+
a := uuid.MustParse("c7c6686d-a93c-4df2-bef9-5f837e9a33d5")
16+
b := uuid.MustParse("8f3b3e0b-2c3f-46a5-a365-fd5b62bd8818")
17+
18+
tests := []struct {
19+
name string
20+
input string
21+
want []uuid.UUID
22+
wantErr string
23+
}{
24+
{
25+
name: "EmptyString",
26+
input: "",
27+
want: []uuid.UUID{},
28+
},
29+
{
30+
name: "JSONNull",
31+
input: "null",
32+
want: []uuid.UUID{},
33+
},
34+
{
35+
name: "WhitespaceOnly",
36+
input: " \n\t ",
37+
want: []uuid.UUID{},
38+
},
39+
{
40+
name: "ValidUUIDs",
41+
input: `["c7c6686d-a93c-4df2-bef9-5f837e9a33d5","8f3b3e0b-2c3f-46a5-a365-fd5b62bd8818"]`,
42+
want: []uuid.UUID{a, b},
43+
},
44+
{
45+
name: "InvalidJSON",
46+
input: "not json at all",
47+
wantErr: "unmarshal uuid list",
48+
},
49+
{
50+
name: "InvalidUUID",
51+
input: `["not-a-uuid"]`,
52+
wantErr: "parse uuid",
53+
},
54+
}
55+
56+
for _, tt := range tests {
57+
t.Run(tt.name, func(t *testing.T) {
58+
t.Parallel()
59+
got, err := xjson.ParseUUIDList(tt.input)
60+
if tt.wantErr != "" {
61+
require.Error(t, err)
62+
require.Contains(t, err.Error(), tt.wantErr)
63+
return
64+
}
65+
require.NoError(t, err)
66+
require.NotNil(t, got)
67+
require.Equal(t, tt.want, got)
68+
})
69+
}
70+
}

‎coderd/x/chatd/chatd.go‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828
"github.com/coder/coder/v2/coderd/database/pubsub"
2929
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
3030
"github.com/coder/coder/v2/coderd/util/ptr"
31+
"github.com/coder/coder/v2/coderd/util/xjson"
3132
"github.com/coder/coder/v2/coderd/webpush"
3233
"github.com/coder/coder/v2/coderd/workspacestats"
3334
"github.com/coder/coder/v2/coderd/x/chatd/chatcost"
@@ -121,6 +122,36 @@ type Server struct {
121122
chatHeartbeatInterval time.Duration
122123
}
123124

125+
// chatTemplateAllowlist returns the deployment-wide template
126+
// allowlist as a set of permitted template IDs. The callback
127+
// signature matches what the chat tools expect. When the
128+
// allowlist is empty or cannot be loaded the function returns
129+
// nil, which the tools interpret as "all templates allowed".
130+
func (p *Server) chatTemplateAllowlist() map[uuid.UUID]bool {
131+
//nolint:gocritic // AsChatd provides narrowly-scoped daemon
132+
// access for reading deployment config.
133+
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
134+
defer cancel()
135+
//nolint:gocritic // AsChatd provides narrowly-scoped read
136+
// access to deployment config (the template allowlist).
137+
ctx = dbauthz.AsChatd(ctx)
138+
raw, err := p.db.GetChatTemplateAllowlist(ctx)
139+
if err != nil {
140+
p.logger.Warn(ctx, "failed to load chat template allowlist", slog.Error(err))
141+
return nil
142+
}
143+
ids, err := xjson.ParseUUIDList(raw)
144+
if err != nil {
145+
p.logger.Warn(ctx, "failed to parse chat template allowlist", slog.Error(err))
146+
return nil
147+
}
148+
m := make(map[uuid.UUID]bool, len(ids))
149+
for _, id := range ids {
150+
m[id] = true
151+
}
152+
return m
153+
}
154+
124155
type turnWorkspaceContext struct {
125156
server *Server
126157
chatStateMu *sync.Mutex
@@ -3413,12 +3444,14 @@ func (p *Server) runChat(
34133444
// Workspace provisioning tools.
34143445
tools = append(tools,
34153446
chattool.ListTemplates(chattool.ListTemplatesOptions{
3416-
DB: p.db,
3417-
OwnerID: chat.OwnerID,
3447+
DB: p.db,
3448+
OwnerID: chat.OwnerID,
3449+
AllowedTemplateIDs: p.chatTemplateAllowlist,
34183450
}),
34193451
chattool.ReadTemplate(chattool.ReadTemplateOptions{
3420-
DB: p.db,
3421-
OwnerID: chat.OwnerID,
3452+
DB: p.db,
3453+
OwnerID: chat.OwnerID,
3454+
AllowedTemplateIDs: p.chatTemplateAllowlist,
34223455
}),
34233456
chattool.CreateWorkspace(chattool.CreateWorkspaceOptions{
34243457
DB: p.db,
@@ -3429,6 +3462,7 @@ func (p *Server) runChat(
34293462
AgentInactiveDisconnectTimeout: p.agentInactiveDisconnectTimeout,
34303463
WorkspaceMu: &workspaceMu,
34313464
Logger: p.logger,
3465+
AllowedTemplateIDs: p.chatTemplateAllowlist,
34323466
}),
34333467
chattool.StartWorkspace(chattool.StartWorkspaceOptions{
34343468
DB: p.db,

‎coderd/x/chatd/chatd_test.go‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3685,3 +3685,116 @@ func TestMCPServerToolInvocation(t *testing.T) {
36853685
require.True(t, foundToolMessage,
36863686
"MCP tool result should be persisted as a tool message in the database")
36873687
}
3688+
3689+
func TestChatTemplateAllowlistEnforcement(t *testing.T) {
3690+
t.Parallel()
3691+
3692+
ctx := testutil.Context(t, testutil.WaitLong)
3693+
db, ps := dbtestutil.NewDB(t)
3694+
3695+
// Set up a mock OpenAI server. The first streaming call triggers
3696+
// list_templates; subsequent calls respond with text.
3697+
var callCount atomic.Int32
3698+
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
3699+
if !req.Stream {
3700+
return chattest.OpenAINonStreamingResponse("title")
3701+
}
3702+
if callCount.Add(1) == 1 {
3703+
return chattest.OpenAIStreamingResponse(
3704+
chattest.OpenAIToolCallChunk("list_templates", `{}`),
3705+
)
3706+
}
3707+
return chattest.OpenAIStreamingResponse(
3708+
chattest.OpenAITextChunks("Here are the templates.")...,
3709+
)
3710+
})
3711+
3712+
user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
3713+
3714+
// Create two templates the user can see.
3715+
org := dbgen.Organization(t, db, database.Organization{})
3716+
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
3717+
UserID: user.ID,
3718+
OrganizationID: org.ID,
3719+
})
3720+
tplAllowed := dbgen.Template(t, db, database.Template{
3721+
OrganizationID: org.ID,
3722+
CreatedBy: user.ID,
3723+
Name: "allowed-template",
3724+
})
3725+
tplBlocked := dbgen.Template(t, db, database.Template{
3726+
OrganizationID: org.ID,
3727+
CreatedBy: user.ID,
3728+
Name: "blocked-template",
3729+
})
3730+
3731+
// Set the allowlist to only tplAllowed.
3732+
allowlistJSON, err := json.Marshal([]string{tplAllowed.ID.String()})
3733+
require.NoError(t, err)
3734+
err = db.UpsertChatTemplateAllowlist(dbauthz.AsSystemRestricted(ctx), string(allowlistJSON))
3735+
require.NoError(t, err)
3736+
3737+
server := newActiveTestServer(t, db, ps)
3738+
3739+
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
3740+
OwnerID: user.ID,
3741+
Title: "allowlist-test",
3742+
ModelConfigID: model.ID,
3743+
InitialUserContent: []codersdk.ChatMessagePart{
3744+
codersdk.ChatMessageText("List templates"),
3745+
},
3746+
})
3747+
require.NoError(t, err)
3748+
3749+
// Wait for the chat to finish processing.
3750+
var chatResult database.Chat
3751+
require.Eventually(t, func() bool {
3752+
got, getErr := db.GetChatByID(ctx, chat.ID)
3753+
if getErr != nil {
3754+
return false
3755+
}
3756+
chatResult = got
3757+
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
3758+
}, testutil.WaitLong, testutil.IntervalFast)
3759+
3760+
if chatResult.Status == database.ChatStatusError {
3761+
require.FailNowf(t, "chat run failed", "last_error=%q", chatResult.LastError.String)
3762+
}
3763+
3764+
// Find the list_templates tool result in the persisted messages.
3765+
var toolResult string
3766+
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
3767+
messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
3768+
ChatID: chat.ID,
3769+
AfterID: 0,
3770+
})
3771+
if dbErr != nil {
3772+
return false
3773+
}
3774+
for _, msg := range messages {
3775+
if msg.Role != database.ChatMessageRoleTool {
3776+
continue
3777+
}
3778+
parts, parseErr := chatprompt.ParseContent(msg)
3779+
if parseErr != nil {
3780+
continue
3781+
}
3782+
for _, part := range parts {
3783+
if part.Type == codersdk.ChatMessagePartTypeToolResult &&
3784+
part.ToolName == "list_templates" {
3785+
toolResult = string(part.Result)
3786+
return true
3787+
}
3788+
}
3789+
}
3790+
return false
3791+
}, testutil.IntervalFast)
3792+
3793+
require.NotEmpty(t, toolResult, "list_templates tool result should be persisted")
3794+
3795+
// The result should contain only the allowed template.
3796+
require.Contains(t, toolResult, tplAllowed.ID.String(),
3797+
"allowed template should appear in list_templates result")
3798+
require.NotContains(t, toolResult, tplBlocked.ID.String(),
3799+
"blocked template should NOT appear in list_templates result")
3800+
}

0 commit comments

Comments
 (0)