完善 MCP 工具

2026-06-11 08:21:09 +08:00
parent c0f15050d0
commit 1415d3b43b
20 changed files with 3722 additions and 61 deletions
@@ -0,0 +1,407 @@
+package mcp
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/yan1h/agent-coding-workflow/internal/chat"
+	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
+	"github.com/yan1h/agent-coding-workflow/internal/user"
+)
+
+// ===== fakes =====
+
+// fakeConvSvc implements chat.ConversationService.
+type fakeConvSvc struct {
+	items []chat.Conversation
+	msgs  []chat.Message
+}
+
+func (f *fakeConvSvc) Create(_ context.Context, userID uuid.UUID, in chat.CreateConversationInput) (*chat.Conversation, error) {
+	cv := chat.Conversation{
+		ID: uuid.New(), UserID: userID, ModelID: in.ModelID,
+		ProjectID: in.ProjectID, WorkspaceID: in.WorkspaceID,
+		IssueID: in.IssueID, RequirementID: in.RequirementID,
+		CreatedAt: time.Now(), UpdatedAt: time.Now(),
+	}
+	f.items = append(f.items, cv)
+	return &cv, nil
+}
+func (f *fakeConvSvc) List(_ context.Context, _ uuid.UUID, filter chat.ConversationFilter) ([]chat.Conversation, error) {
+	out := make([]chat.Conversation, 0, len(f.items))
+	for _, cv := range f.items {
+		if filter.ProjectID != nil && (cv.ProjectID == nil || *cv.ProjectID != *filter.ProjectID) {
+			continue
+		}
+		out = append(out, cv)
+	}
+	return out, nil
+}
+func (f *fakeConvSvc) Get(_ context.Context, _ uuid.UUID, id uuid.UUID, _ int) (*chat.Conversation, []chat.Message, error) {
+	for i := range f.items {
+		if f.items[i].ID == id {
+			return &f.items[i], f.msgs, nil
+		}
+	}
+	return nil, nil, errs.New(errs.CodeNotFound, "conversation")
+}
+func (f *fakeConvSvc) UpdateTitle(_ context.Context, _ uuid.UUID, id uuid.UUID, title string) error {
+	for i := range f.items {
+		if f.items[i].ID == id {
+			f.items[i].Title = &title
+			return nil
+		}
+	}
+	return errs.New(errs.CodeNotFound, "conversation")
+}
+func (f *fakeConvSvc) SoftDelete(_ context.Context, _ uuid.UUID, id uuid.UUID) error {
+	kept := f.items[:0]
+	found := false
+	for _, cv := range f.items {
+		if cv.ID == id {
+			found = true
+			continue
+		}
+		kept = append(kept, cv)
+	}
+	f.items = kept
+	if !found {
+		return errs.New(errs.CodeNotFound, "conversation")
+	}
+	return nil
+}
+
+// fakeMsgSvc implements chat.MessageService.
+type fakeMsgSvc struct {
+	sentContent string
+	msgs        []chat.Message
+}
+
+func (f *fakeMsgSvc) Send(_ context.Context, _ uuid.UUID, _ uuid.UUID, content string, _ []uuid.UUID) (int64, int64, error) {
+	f.sentContent = content
+	return 1, 2, nil
+}
+func (f *fakeMsgSvc) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan chat.SSEEvent, error) {
+	return nil, errs.New(errs.CodeInternal, "not implemented")
+}
+func (f *fakeMsgSvc) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil }
+func (f *fakeMsgSvc) Retry(_ context.Context, _ uuid.UUID, _ int64) (int64, uuid.UUID, error) {
+	return 0, uuid.Nil, errs.New(errs.CodeInternal, "not implemented")
+}
+func (f *fakeMsgSvc) ListMessages(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ *int64, _ int) ([]chat.Message, error) {
+	return f.msgs, nil
+}
+
+// fakeTplSvc implements chat.TemplateService.
+type fakeTplSvc struct {
+	items []chat.PromptTemplate
+}
+
+func (f *fakeTplSvc) List(_ context.Context, _ uuid.UUID) ([]chat.PromptTemplate, error) {
+	return f.items, nil
+}
+func (f *fakeTplSvc) Create(_ context.Context, userID uuid.UUID, name, content string, scope chat.TemplateScope) (*chat.PromptTemplate, error) {
+	t := chat.PromptTemplate{ID: uuid.New(), Name: name, Content: content, Scope: scope}
+	if scope == chat.TemplateScopeUser {
+		uid := userID
+		t.OwnerID = &uid
+	}
+	f.items = append(f.items, t)
+	return &t, nil
+}
+func (f *fakeTplSvc) Update(_ context.Context, _ uuid.UUID, _ bool, id uuid.UUID, name, content string) error {
+	for i := range f.items {
+		if f.items[i].ID == id {
+			f.items[i].Name = name
+			f.items[i].Content = content
+			return nil
+		}
+	}
+	return errs.New(errs.CodeChatTemplateNotFound, "template not found")
+}
+func (f *fakeTplSvc) Delete(_ context.Context, _ uuid.UUID, _ bool, id uuid.UUID) error {
+	kept := f.items[:0]
+	for _, t := range f.items {
+		if t.ID != id {
+			kept = append(kept, t)
+		}
+	}
+	f.items = kept
+	return nil
+}
+
+// fakeEndpointSvc implements chat.EndpointService（仅 ListModels 有效）。
+type fakeEndpointSvc struct {
+	models []chat.LLMModel
+}
+
+func (f *fakeEndpointSvc) ListEndpoints(_ context.Context) ([]chat.LLMEndpoint, error) {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) CreateEndpoint(_ context.Context, _ uuid.UUID, _ chat.CreateEndpointInput) (*chat.LLMEndpoint, error) {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) UpdateEndpoint(_ context.Context, _ uuid.UUID, _ chat.UpdateEndpointInput) error {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) DeleteEndpoint(_ context.Context, _ uuid.UUID) error {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) TestEndpoint(_ context.Context, _ uuid.UUID) error {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) ListModels(_ context.Context, onlyEnabled bool) ([]chat.LLMModel, error) {
+	if !onlyEnabled {
+		return f.models, nil
+	}
+	out := make([]chat.LLMModel, 0, len(f.models))
+	for _, m := range f.models {
+		if m.Enabled {
+			out = append(out, m)
+		}
+	}
+	return out, nil
+}
+func (f *fakeEndpointSvc) CreateModel(_ context.Context, _ chat.CreateModelInput) (*chat.LLMModel, error) {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) UpdateModel(_ context.Context, _ uuid.UUID, _ chat.UpdateModelInput) error {
+	panic("not implemented")
+}
+func (f *fakeEndpointSvc) DeleteModel(_ context.Context, _ uuid.UUID) error {
+	panic("not implemented")
+}
+
+// ===== helpers =====
+
+// testUID2 是非 admin 用户，用于权限边界测试。
+var testUID2 = uuid.MustParse("dddddddd-dddd-dddd-dddd-dddddddddddd")
+
+func chatTestDeps() (ServerDeps, *fakeConvSvc, *fakeMsgSvc, *fakeTplSvc) {
+	deps := testDeps()
+	deps.Caller = NewCallerResolver(&fakeUserSvc{users: map[uuid.UUID]*user.User{
+		testUID:  {ID: testUID, IsAdmin: true},
+		testUID2: {ID: testUID2, IsAdmin: false},
+	}})
+	convSvc := &fakeConvSvc{}
+	msgSvc := &fakeMsgSvc{}
+	tplSvc := &fakeTplSvc{}
+	deps.Conversations = convSvc
+	deps.Messages = msgSvc
+	deps.Templates = tplSvc
+	deps.Endpoints = &fakeEndpointSvc{}
+	return deps, convSvc, msgSvc, tplSvc
+}
+
+func ctxWithAuthUser(uid uuid.UUID, scope Scope) context.Context {
+	return context.WithValue(context.Background(), AuthContextKey, &AuthSession{
+		UserID:  uid.String(),
+		TokenID: uuid.NewString(),
+		Scope:   scope,
+	})
+}
+
+// ===== tests =====
+
+func TestCreateAndGetConversation(t *testing.T) {
+	deps, _, _, _ := chatTestDeps()
+	modelID := uuid.New()
+
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "create_conversation", map[string]any{
+		"model_id": modelID.String(), "project_slug": "test-project",
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	var createOut conversationCreateOut
+	unmarshalStructured(t, result, &createOut)
+	assert.True(t, createOut.OK)
+	assert.Equal(t, modelID.String(), createOut.Conversation.ModelID)
+	assert.Equal(t, testPID.String(), createOut.Conversation.ProjectID)
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "get_conversation", map[string]any{
+		"conversation_id": createOut.Conversation.ID,
+	})
+	require.NoError(t, err)
+	var getOut getConversationOut
+	unmarshalStructured(t, result, &getOut)
+	assert.Equal(t, createOut.Conversation.ID, getOut.Conversation.ID)
+}
+
+func TestCreateConversation_NumberNeedsProjectSlug(t *testing.T) {
+	deps, _, _, _ := chatTestDeps()
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "create_conversation", map[string]any{
+		"model_id": uuid.New().String(), "issue_number": float64(1),
+	})
+	require.NoError(t, err)
+	assert.True(t, result.IsError)
+}
+
+func TestListConversations_FilterByProject(t *testing.T) {
+	deps, convSvc, _, _ := chatTestDeps()
+	pid := testPID
+	convSvc.items = []chat.Conversation{
+		{ID: uuid.New(), ModelID: uuid.New(), ProjectID: &pid, CreatedAt: time.Now(), UpdatedAt: time.Now()},
+		{ID: uuid.New(), ModelID: uuid.New(), CreatedAt: time.Now(), UpdatedAt: time.Now()},
+	}
+
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "list_conversations", map[string]any{
+		"project_slug": "test-project",
+	})
+	require.NoError(t, err)
+	var out listConversationsOut
+	unmarshalStructured(t, result, &out)
+	require.Len(t, out.Conversations, 1)
+	assert.Equal(t, pid.String(), out.Conversations[0].ProjectID)
+}
+
+func TestUpdateTitleAndDeleteConversation(t *testing.T) {
+	deps, convSvc, _, _ := chatTestDeps()
+	cv := chat.Conversation{ID: uuid.New(), ModelID: uuid.New(), CreatedAt: time.Now(), UpdatedAt: time.Now()}
+	convSvc.items = []chat.Conversation{cv}
+
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "update_conversation_title", map[string]any{
+		"conversation_id": cv.ID.String(), "title": "New Title",
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	require.NotNil(t, convSvc.items[0].Title)
+	assert.Equal(t, "New Title", *convSvc.items[0].Title)
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "delete_conversation", map[string]any{
+		"conversation_id": cv.ID.String(),
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	assert.Empty(t, convSvc.items)
+}
+
+func TestSendMessage(t *testing.T) {
+	deps, _, msgSvc, _ := chatTestDeps()
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "send_message", map[string]any{
+		"conversation_id": uuid.New().String(), "content": "hello",
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	var out sendMessageOut
+	unmarshalStructured(t, result, &out)
+	assert.Equal(t, int64(1), out.UserMessageID)
+	assert.Equal(t, int64(2), out.AssistantMessageID)
+	assert.Equal(t, "hello", msgSvc.sentContent)
+}
+
+func TestPromptTemplateCRUD(t *testing.T) {
+	deps, _, _, tplSvc := chatTestDeps()
+
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "create_prompt_template", map[string]any{
+		"name": "tpl", "content": "You are helpful.",
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	var createOut templateCreateOut
+	unmarshalStructured(t, result, &createOut)
+	assert.Equal(t, "user", createOut.Template.Scope, "scope defaults to user")
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "update_prompt_template", map[string]any{
+		"template_id": createOut.Template.ID, "name": "tpl2", "content": "Updated.",
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	assert.Equal(t, "tpl2", tplSvc.items[0].Name)
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "list_prompt_templates", map[string]any{})
+	require.NoError(t, err)
+	var listOut listPromptTemplatesOut
+	unmarshalStructured(t, result, &listOut)
+	require.Len(t, listOut.Templates, 1)
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "delete_prompt_template", map[string]any{
+		"template_id": createOut.Template.ID,
+	})
+	require.NoError(t, err)
+	require.False(t, result.IsError)
+	assert.Empty(t, tplSvc.items)
+}
+
+func TestCreateSystemTemplate_AdminOnly(t *testing.T) {
+	deps, _, _, tplSvc := chatTestDeps()
+
+	// 非 admin → 拒绝
+	result, err := callTool(ctxWithAuthUser(testUID2, Scope{}), deps, "create_prompt_template", map[string]any{
+		"name": "sys", "content": "x", "scope": "system",
+	})
+	require.NoError(t, err)
+	assert.True(t, result.IsError)
+	assert.Empty(t, tplSvc.items)
+
+	// admin → 允许
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "create_prompt_template", map[string]any{
+		"name": "sys", "content": "x", "scope": "system",
+	})
+	require.NoError(t, err)
+	assert.False(t, result.IsError)
+	require.Len(t, tplSvc.items, 1)
+	assert.Equal(t, chat.TemplateScopeSystem, tplSvc.items[0].Scope)
+}
+
+func TestUpdateDeletePromptTemplate_PassesAdminFlag(t *testing.T) {
+	deps, _, _, _ := chatTestDeps()
+	rec := &adminFlagRecordingTplSvc{}
+	deps.Templates = rec
+	tid := uuid.New()
+
+	// admin caller（testUID）→ isAdmin=true 透传到 service
+	_, err := callTool(ctxWithAuth(Scope{}), deps, "update_prompt_template", map[string]any{
+		"template_id": tid.String(), "name": "n", "content": "c",
+	})
+	require.NoError(t, err)
+	assert.True(t, rec.lastIsAdmin)
+
+	// 非 admin caller（testUID2）→ isAdmin=false
+	_, err = callTool(ctxWithAuthUser(testUID2, Scope{}), deps, "delete_prompt_template", map[string]any{
+		"template_id": tid.String(),
+	})
+	require.NoError(t, err)
+	assert.False(t, rec.lastIsAdmin)
+}
+
+// adminFlagRecordingTplSvc 记录 Update/Delete 收到的 isAdmin 标志。
+type adminFlagRecordingTplSvc struct {
+	fakeTplSvc
+	lastIsAdmin bool
+}
+
+func (f *adminFlagRecordingTplSvc) Update(_ context.Context, _ uuid.UUID, isAdmin bool, _ uuid.UUID, _, _ string) error {
+	f.lastIsAdmin = isAdmin
+	return nil
+}
+func (f *adminFlagRecordingTplSvc) Delete(_ context.Context, _ uuid.UUID, isAdmin bool, _ uuid.UUID) error {
+	f.lastIsAdmin = isAdmin
+	return nil
+}
+
+func TestListModels(t *testing.T) {
+	deps, _, _, _ := chatTestDeps()
+	deps.Endpoints = &fakeEndpointSvc{models: []chat.LLMModel{
+		{ID: uuid.New(), ModelID: "m-on", DisplayName: "On", Enabled: true},
+		{ID: uuid.New(), ModelID: "m-off", DisplayName: "Off", Enabled: false},
+	}}
+
+	result, err := callTool(ctxWithAuth(Scope{}), deps, "list_models", map[string]any{})
+	require.NoError(t, err)
+	var out listModelsOut
+	unmarshalStructured(t, result, &out)
+	require.Len(t, out.Models, 1)
+	assert.Equal(t, "m-on", out.Models[0].ModelID)
+
+	result, err = callTool(ctxWithAuth(Scope{}), deps, "list_models", map[string]any{
+		"include_disabled": true,
+	})
+	require.NoError(t, err)
+	unmarshalStructured(t, result, &out)
+	assert.Len(t, out.Models, 2)
+}