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.
This commit is contained in:
Cian Johnston
2026-03-25 22:15:27 +00:00
committed by GitHub
parent dab4e6f0a4
commit 7a9d57cd87
5 changed files with 262 additions and 23 deletions
+38 -4
View File
@@ -28,6 +28,7 @@ import (
"github.com/coder/coder/v2/coderd/database/pubsub"
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/coderd/util/xjson"
"github.com/coder/coder/v2/coderd/webpush"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/coderd/x/chatd/chatcost"
@@ -121,6 +122,36 @@ type Server struct {
chatHeartbeatInterval time.Duration
}
// chatTemplateAllowlist returns the deployment-wide template
// allowlist as a set of permitted template IDs. The callback
// signature matches what the chat tools expect. When the
// allowlist is empty or cannot be loaded the function returns
// nil, which the tools interpret as "all templates allowed".
func (p *Server) chatTemplateAllowlist() map[uuid.UUID]bool {
//nolint:gocritic // AsChatd provides narrowly-scoped daemon
// access for reading deployment config.
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
//nolint:gocritic // AsChatd provides narrowly-scoped read
// access to deployment config (the template allowlist).
ctx = dbauthz.AsChatd(ctx)
raw, err := p.db.GetChatTemplateAllowlist(ctx)
if err != nil {
p.logger.Warn(ctx, "failed to load chat template allowlist", slog.Error(err))
return nil
}
ids, err := xjson.ParseUUIDList(raw)
if err != nil {
p.logger.Warn(ctx, "failed to parse chat template allowlist", slog.Error(err))
return nil
}
m := make(map[uuid.UUID]bool, len(ids))
for _, id := range ids {
m[id] = true
}
return m
}
type turnWorkspaceContext struct {
server *Server
chatStateMu *sync.Mutex
@@ -3413,12 +3444,14 @@ func (p *Server) runChat(
// Workspace provisioning tools.
tools = append(tools,
chattool.ListTemplates(chattool.ListTemplatesOptions{
DB: p.db,
OwnerID: chat.OwnerID,
DB: p.db,
OwnerID: chat.OwnerID,
AllowedTemplateIDs: p.chatTemplateAllowlist,
}),
chattool.ReadTemplate(chattool.ReadTemplateOptions{
DB: p.db,
OwnerID: chat.OwnerID,
DB: p.db,
OwnerID: chat.OwnerID,
AllowedTemplateIDs: p.chatTemplateAllowlist,
}),
chattool.CreateWorkspace(chattool.CreateWorkspaceOptions{
DB: p.db,
@@ -3429,6 +3462,7 @@ func (p *Server) runChat(
AgentInactiveDisconnectTimeout: p.agentInactiveDisconnectTimeout,
WorkspaceMu: &workspaceMu,
Logger: p.logger,
AllowedTemplateIDs: p.chatTemplateAllowlist,
}),
chattool.StartWorkspace(chattool.StartWorkspaceOptions{
DB: p.db,
+113
View File
@@ -3685,3 +3685,116 @@ func TestMCPServerToolInvocation(t *testing.T) {
require.True(t, foundToolMessage,
"MCP tool result should be persisted as a tool message in the database")
}
func TestChatTemplateAllowlistEnforcement(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, ps := dbtestutil.NewDB(t)
// Set up a mock OpenAI server. The first streaming call triggers
// list_templates; subsequent calls respond with text.
var callCount atomic.Int32
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
if !req.Stream {
return chattest.OpenAINonStreamingResponse("title")
}
if callCount.Add(1) == 1 {
return chattest.OpenAIStreamingResponse(
chattest.OpenAIToolCallChunk("list_templates", `{}`),
)
}
return chattest.OpenAIStreamingResponse(
chattest.OpenAITextChunks("Here are the templates.")...,
)
})
user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
// Create two templates the user can see.
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
tplAllowed := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
Name: "allowed-template",
})
tplBlocked := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
Name: "blocked-template",
})
// Set the allowlist to only tplAllowed.
allowlistJSON, err := json.Marshal([]string{tplAllowed.ID.String()})
require.NoError(t, err)
err = db.UpsertChatTemplateAllowlist(dbauthz.AsSystemRestricted(ctx), string(allowlistJSON))
require.NoError(t, err)
server := newActiveTestServer(t, db, ps)
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
OwnerID: user.ID,
Title: "allowlist-test",
ModelConfigID: model.ID,
InitialUserContent: []codersdk.ChatMessagePart{
codersdk.ChatMessageText("List templates"),
},
})
require.NoError(t, err)
// Wait for the chat to finish processing.
var chatResult database.Chat
require.Eventually(t, func() bool {
got, getErr := db.GetChatByID(ctx, chat.ID)
if getErr != nil {
return false
}
chatResult = got
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
}, testutil.WaitLong, testutil.IntervalFast)
if chatResult.Status == database.ChatStatusError {
require.FailNowf(t, "chat run failed", "last_error=%q", chatResult.LastError.String)
}
// Find the list_templates tool result in the persisted messages.
var toolResult string
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
ChatID: chat.ID,
AfterID: 0,
})
if dbErr != nil {
return false
}
for _, msg := range messages {
if msg.Role != database.ChatMessageRoleTool {
continue
}
parts, parseErr := chatprompt.ParseContent(msg)
if parseErr != nil {
continue
}
for _, part := range parts {
if part.Type == codersdk.ChatMessagePartTypeToolResult &&
part.ToolName == "list_templates" {
toolResult = string(part.Result)
return true
}
}
}
return false
}, testutil.IntervalFast)
require.NotEmpty(t, toolResult, "list_templates tool result should be persisted")
// The result should contain only the allowed template.
require.Contains(t, toolResult, tplAllowed.ID.String(),
"allowed template should appear in list_templates result")
require.NotContains(t, toolResult, tplBlocked.ID.String(),
"blocked template should NOT appear in list_templates result")
}