agentic-coding-workflow/internal/mcp/tools_memory.go

package mcp

import (
	"context"

	"github.com/google/uuid"
	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	"github.com/yan1h/agent-coding-workflow/internal/project"
	"github.com/yan1h/agent-coding-workflow/internal/projectmemory"
)

// registerMemoryTools registers the project_memory tools. When the Memory service
// is nil the tools are not registered.
func registerMemoryTools(srv *mcpsdk.Server, deps ServerDeps) {
	if deps.Memory == nil {
		return
	}
	registerMemoryWrite(srv, deps)
	registerMemorySearch(srv, deps)
	registerMemoryList(srv, deps)
}

// resolveMemoryScope resolves (projectID, workspaceID?) from either a
// workspace_id or a project_slug, enforcing project scope. Exactly one of the two
// identifiers must be supplied for write/search; list requires project_slug.
func resolveMemoryScope(ctx context.Context, deps ServerDeps, c project.Caller, workspaceID, projectSlug string) (projectID uuid.UUID, wsID *uuid.UUID, err error) {
	switch {
	case workspaceID != "":
		w, _, werr := scopedWorkspace(ctx, deps, c, workspaceID)
		if werr != nil {
			return uuid.Nil, nil, werr
		}
		id := w.ID
		return w.ProjectID, &id, nil
	case projectSlug != "":
		p, perr := deps.Projects.Get(ctx, c, projectSlug)
		if perr != nil {
			return uuid.Nil, nil, perr
		}
		if cerr := CheckProjectFromCtx(ctx, p.ID); cerr != nil {
			return uuid.Nil, nil, cerr
		}
		return p.ID, nil, nil
	default:
		return uuid.Nil, nil, errs.New(errs.CodeInvalidInput, "project_slug or workspace_id required")
	}
}

type memoryEntryDTO struct {
	ID          string   `json:"id"`
	Kind        string   `json:"kind"`
	Title       string   `json:"title"`
	Body        string   `json:"body"`
	Tags        []string `json:"tags"`
	Source      string   `json:"source"`
	WorkspaceID string   `json:"workspace_id,omitempty"`
	CreatedAt   string   `json:"created_at"`
}

func memoryEntryDTOFrom(e projectmemory.Entry) memoryEntryDTO {
	d := memoryEntryDTO{
		ID: e.ID.String(), Kind: e.Kind, Title: e.Title, Body: e.Body,
		Tags: e.Tags, Source: e.Source,
		CreatedAt: e.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
	}
	if e.WorkspaceID != nil {
		d.WorkspaceID = e.WorkspaceID.String()
	}
	if d.Tags == nil {
		d.Tags = []string{}
	}
	return d
}

// ===== memory_write =====

type memoryWriteIn struct {
	ProjectSlug string   `json:"project_slug,omitempty" jsonschema:"project slug (or supply workspace_id)"`
	WorkspaceID string   `json:"workspace_id,omitempty" jsonschema:"workspace UUID (scopes the entry to this workspace)"`
	Kind        string   `json:"kind" jsonschema:"one of: decision, convention, file_map, gotcha"`
	Title       string   `json:"title"`
	Body        string   `json:"body"`
	Tags        []string `json:"tags,omitempty"`
}

type memoryWriteOut struct {
	OK bool   `json:"ok"`
	ID string `json:"id"`
}

func registerMemoryWrite(srv *mcpsdk.Server, deps ServerDeps) {
	mcpsdk.AddTool(srv,
		&mcpsdk.Tool{Name: "memory_write", Description: "Record a typed project memory entry (decision/convention/file_map/gotcha)."},
		func(ctx context.Context, _ *mcpsdk.CallToolRequest, in memoryWriteIn) (*mcpsdk.CallToolResult, memoryWriteOut, error) {
			if err := CheckToolFromCtx(ctx, "memory_write"); err != nil {
				return nil, memoryWriteOut{}, err
			}
			c, err := deps.Caller.Resolve(ctx)
			if err != nil {
				return nil, memoryWriteOut{}, err
			}
			projectID, wsID, err := resolveMemoryScope(ctx, deps, c, in.WorkspaceID, in.ProjectSlug)
			if err != nil {
				return nil, memoryWriteOut{}, err
			}
			uid := c.UserID
			entry, err := deps.Memory.Write(ctx, projectmemory.WriteInput{
				ProjectID:   projectID,
				WorkspaceID: wsID,
				Kind:        in.Kind,
				Title:       in.Title,
				Body:        in.Body,
				Tags:        in.Tags,
				Source:      projectmemory.SourceAgent,
				CreatedBy:   &uid,
			})
			if err != nil {
				return nil, memoryWriteOut{}, err
			}
			return nil, memoryWriteOut{OK: true, ID: entry.ID.String()}, nil
		})
}

// ===== memory_search =====

type memorySearchIn struct {
	ProjectSlug string `json:"project_slug,omitempty" jsonschema:"project slug (or supply workspace_id)"`
	WorkspaceID string `json:"workspace_id,omitempty" jsonschema:"workspace UUID"`
	Query       string `json:"query" jsonschema:"search text"`
	Kind        string `json:"kind,omitempty" jsonschema:"restrict to a kind"`
	TopK        int    `json:"top_k,omitempty"`
}

type memorySearchOut struct {
	Results []memoryEntryDTO `json:"results"`
}

func registerMemorySearch(srv *mcpsdk.Server, deps ServerDeps) {
	mcpsdk.AddTool(srv,
		&mcpsdk.Tool{Name: "memory_search", Description: "Search project memory (vector when embeddings available, keyword/tag fallback otherwise)."},
		func(ctx context.Context, _ *mcpsdk.CallToolRequest, in memorySearchIn) (*mcpsdk.CallToolResult, memorySearchOut, error) {
			if err := CheckToolFromCtx(ctx, "memory_search"); err != nil {
				return nil, memorySearchOut{}, err
			}
			c, err := deps.Caller.Resolve(ctx)
			if err != nil {
				return nil, memorySearchOut{}, err
			}
			projectID, _, err := resolveMemoryScope(ctx, deps, c, in.WorkspaceID, in.ProjectSlug)
			if err != nil {
				return nil, memorySearchOut{}, err
			}
			entries, err := deps.Memory.Search(ctx, projectmemory.Query{
				ProjectID: projectID, Text: in.Query, Kind: in.Kind, TopK: in.TopK,
			})
			if err != nil {
				return nil, memorySearchOut{}, err
			}
			out := memorySearchOut{Results: make([]memoryEntryDTO, 0, len(entries))}
			for _, e := range entries {
				out.Results = append(out.Results, memoryEntryDTOFrom(e))
			}
			return nil, out, nil
		})
}

// ===== memory_list =====

type memoryListIn struct {
	ProjectSlug string `json:"project_slug" jsonschema:"project slug"`
	Kind        string `json:"kind,omitempty" jsonschema:"restrict to a kind"`
}

type memoryListOut struct {
	Entries []memoryEntryDTO `json:"entries"`
}

func registerMemoryList(srv *mcpsdk.Server, deps ServerDeps) {
	mcpsdk.AddTool(srv,
		&mcpsdk.Tool{Name: "memory_list", Description: "List project memory entries, optionally filtered by kind."},
		func(ctx context.Context, _ *mcpsdk.CallToolRequest, in memoryListIn) (*mcpsdk.CallToolResult, memoryListOut, error) {
			if err := CheckToolFromCtx(ctx, "memory_list"); err != nil {
				return nil, memoryListOut{}, err
			}
			c, err := deps.Caller.Resolve(ctx)
			if err != nil {
				return nil, memoryListOut{}, err
			}
			if in.ProjectSlug == "" {
				return nil, memoryListOut{}, errs.New(errs.CodeInvalidInput, "project_slug required")
			}
			p, err := deps.Projects.Get(ctx, c, in.ProjectSlug)
			if err != nil {
				return nil, memoryListOut{}, err
			}
			if err := CheckProjectFromCtx(ctx, p.ID); err != nil {
				return nil, memoryListOut{}, err
			}
			entries, err := deps.Memory.List(ctx, p.ID, in.Kind)
			if err != nil {
				return nil, memoryListOut{}, err
			}
			out := memoryListOut{Entries: make([]memoryEntryDTO, 0, len(entries))}
			for _, e := range entries {
				out.Entries = append(out.Entries, memoryEntryDTOFrom(e))
			}
			return nil, out, nil
		})
}