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 }) }