mirror of
https://github.com/coder/coder.git
synced 2026-10-24 11:43:39 +08:00
Cian Johnston
6194bd6f57
Addresses review findings from #23827 that were added post-merge: - Persisted attachments now store `organizationId`; mismatched orgs pruned on restore - Workspace selection reconciliation: stale IDs from previous orgs dropped via derived `effectiveWorkspaceId` - Org picker uses `permittedOrganizations()` for RBAC-aware filtering - Org picker hidden when user belongs to only one org - Ref-sync `useEffect` replaced with `useEffectEvent` - `CreateWorkspace()` and `ListTemplates()` take `organizationID` and `db` as required function parameters instead of optional struct fields — compiler enforces them, removes scattered nil guards - Cross-org template check in `CreateWorkspace` is now unconditional - `ListTemplates` org-scoping filter now has test coverage - `setupChatInfra` comment fixed; test helpers use params structs instead of positional UUIDs - Enterprise test documents that org admin only sees own chats (handler hardcodes `OwnerID` — future work needs sidebar UI before lifting that restriction) > 🤖
291 lines
10 KiB
Go
291 lines
10 KiB
Go
package chattool_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"charm.land/fantasy"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbgen"
|
|
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
func TestListTemplates_OrganizationFilter(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
user := dbgen.User(t, db, database.User{})
|
|
|
|
orgA := dbgen.Organization(t, db, database.Organization{})
|
|
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
UserID: user.ID,
|
|
OrganizationID: orgA.ID,
|
|
})
|
|
orgB := dbgen.Organization(t, db, database.Organization{})
|
|
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
UserID: user.ID,
|
|
OrganizationID: orgB.ID,
|
|
})
|
|
|
|
tAlpha := dbgen.Template(t, db, database.Template{
|
|
OrganizationID: orgA.ID,
|
|
CreatedBy: user.ID,
|
|
Name: "alpha",
|
|
})
|
|
tBeta := dbgen.Template(t, db, database.Template{
|
|
OrganizationID: orgB.ID,
|
|
CreatedBy: user.ID,
|
|
Name: "beta",
|
|
})
|
|
|
|
t.Run("ScopedToOrgA", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tool := chattool.ListTemplates(orgA.ID, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "org-a", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Len(t, templates, 1)
|
|
m := templates[0].(map[string]any)
|
|
require.Equal(t, tAlpha.ID.String(), m["id"].(string))
|
|
})
|
|
|
|
t.Run("NilOrgReturnsBoth", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
// Pass uuid.Nil to skip org filtering.
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "nil-org", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Len(t, templates, 2)
|
|
})
|
|
|
|
t.Run("ReadTemplate_CrossOrgRejected", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
// Tool scoped to orgA, but requesting a template in orgB.
|
|
tool := chattool.ReadTemplate(orgA.ID, db, chattool.ReadTemplateOptions{
|
|
OwnerID: user.ID,
|
|
})
|
|
|
|
input := `{"template_id":"` + tBeta.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "cross-org", Name: "read_template", Input: input})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "not found")
|
|
})
|
|
|
|
t.Run("ReadTemplate_SameOrgAllowed", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
// Tool scoped to orgA, requesting a template in orgA.
|
|
tool := chattool.ReadTemplate(orgA.ID, db, chattool.ReadTemplateOptions{
|
|
OwnerID: user.ID,
|
|
})
|
|
|
|
input := `{"template_id":"` + tAlpha.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "same-org", Name: "read_template", Input: input})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
tmplInfo := result["template"].(map[string]any)
|
|
require.Equal(t, tAlpha.ID.String(), tmplInfo["id"].(string))
|
|
})
|
|
}
|
|
|
|
//nolint:tparallel,paralleltest // Subtests share a single DB and run sequentially.
|
|
func TestTemplateAllowlistEnforcement(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
db, _ := dbtestutil.NewDB(t)
|
|
|
|
user := dbgen.User(t, db, database.User{})
|
|
org := dbgen.Organization(t, db, database.Organization{})
|
|
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
UserID: user.ID,
|
|
OrganizationID: org.ID,
|
|
})
|
|
|
|
t1 := dbgen.Template(t, db, database.Template{
|
|
OrganizationID: org.ID,
|
|
CreatedBy: user.ID,
|
|
Name: "template-alpha",
|
|
})
|
|
t2 := dbgen.Template(t, db, database.Template{
|
|
OrganizationID: org.ID,
|
|
CreatedBy: user.ID,
|
|
Name: "template-beta",
|
|
})
|
|
|
|
t.Run("ListTemplates", func(t *testing.T) {
|
|
t.Run("NoAllowlist", func(t *testing.T) {
|
|
tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c1", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Len(t, templates, 2)
|
|
})
|
|
|
|
t.Run("EmptyAllowlist", func(t *testing.T) {
|
|
tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{} },
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c2", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Len(t, templates, 2)
|
|
})
|
|
|
|
t.Run("OneMatch", func(t *testing.T) {
|
|
tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c3", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Len(t, templates, 1)
|
|
m := templates[0].(map[string]any)
|
|
require.Equal(t, t1.ID.String(), m["id"].(string))
|
|
})
|
|
|
|
t.Run("NoMatches", func(t *testing.T) {
|
|
tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} },
|
|
})
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c4", Name: "list_templates", Input: "{}"})
|
|
require.NoError(t, err)
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
templates := result["templates"].([]any)
|
|
require.Empty(t, templates)
|
|
})
|
|
})
|
|
|
|
t.Run("ReadTemplate", func(t *testing.T) {
|
|
t.Run("Allowed", func(t *testing.T) {
|
|
tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
|
|
})
|
|
input := `{"template_id":"` + t1.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c5", Name: "read_template", Input: input})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError)
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
tmplInfo := result["template"].(map[string]any)
|
|
require.Equal(t, t1.ID.String(), tmplInfo["id"].(string))
|
|
})
|
|
|
|
t.Run("Disallowed", func(t *testing.T) {
|
|
tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} },
|
|
})
|
|
input := `{"template_id":"` + t2.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c6", Name: "read_template", Input: input})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "not found")
|
|
})
|
|
|
|
t.Run("NoAllowlist", func(t *testing.T) {
|
|
tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{
|
|
OwnerID: user.ID,
|
|
})
|
|
input := `{"template_id":"` + t2.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c7", Name: "read_template", Input: input})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError)
|
|
})
|
|
})
|
|
|
|
t.Run("CreateWorkspace", func(t *testing.T) {
|
|
t.Run("Allowed", func(t *testing.T) {
|
|
createCalled := false
|
|
tool := chattool.CreateWorkspace(org.ID, db, chattool.CreateWorkspaceOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
|
|
|
|
CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
createCalled = true
|
|
return codersdk.Workspace{}, nil
|
|
},
|
|
})
|
|
|
|
input := `{"template_id":"` + t1.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c8a", Name: "create_workspace", Input: input})
|
|
require.NoError(t, err)
|
|
require.True(t, createCalled, "CreateFn should be called for allowed template")
|
|
// We don't assert resp.IsError here because CreateWorkspace
|
|
// does additional work (asOwner, workspace lookup) that
|
|
// depends on full RBAC setup. The key assertion is that
|
|
// the allowlist gate passed and CreateFn was invoked.
|
|
_ = resp
|
|
})
|
|
|
|
t.Run("Disallowed", func(t *testing.T) {
|
|
createCalled := false
|
|
tool := chattool.CreateWorkspace(uuid.Nil, db, chattool.CreateWorkspaceOptions{
|
|
OwnerID: user.ID,
|
|
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} },
|
|
CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
createCalled = true
|
|
t.Fatal("CreateFn should not be called for blocked template")
|
|
return codersdk.Workspace{}, nil
|
|
},
|
|
})
|
|
|
|
input := `{"template_id":"` + t1.ID.String() + `"}`
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c8", Name: "create_workspace", Input: input})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "template not available for chat workspaces")
|
|
require.False(t, createCalled, "CreateFn should not be called for blocked template")
|
|
})
|
|
})
|
|
}
|