You've already forked agentic-coding-workflow
110 lines
3.7 KiB
Go
110 lines
3.7 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
|
|
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/codeindex"
|
|
)
|
|
|
|
// registerCodeIndexTools registers the pgvector code-search tools. When the
|
|
// CodeIndex service is not configured (no embedding endpoint), the tools are not
|
|
// registered at all so they don't appear in the agent's tool list.
|
|
func registerCodeIndexTools(srv *mcpsdk.Server, deps ServerDeps) {
|
|
if deps.CodeIndex == nil {
|
|
return
|
|
}
|
|
registerSearchCode(srv, deps)
|
|
registerReindexWorkspace(srv, deps)
|
|
}
|
|
|
|
// ===== search_code =====
|
|
|
|
type searchCodeIn struct {
|
|
WorkspaceID string `json:"workspace_id" jsonschema:"workspace UUID"`
|
|
Query string `json:"query" jsonschema:"natural-language or code query"`
|
|
TopK int `json:"top_k,omitempty" jsonschema:"max results (capped by server)"`
|
|
PathPrefix string `json:"path_prefix,omitempty" jsonschema:"restrict to a worktree-relative path prefix"`
|
|
Lang string `json:"lang,omitempty" jsonschema:"restrict to a language label (e.g. go, typescript)"`
|
|
}
|
|
|
|
type codeSnippet struct {
|
|
File string `json:"file"`
|
|
StartLine int `json:"start_line"`
|
|
EndLine int `json:"end_line"`
|
|
Snippet string `json:"snippet"`
|
|
Score float32 `json:"score"`
|
|
CommitSHA string `json:"commit_sha"`
|
|
}
|
|
|
|
type searchCodeOut struct {
|
|
Results []codeSnippet `json:"results"`
|
|
}
|
|
|
|
func registerSearchCode(srv *mcpsdk.Server, deps ServerDeps) {
|
|
mcpsdk.AddTool(srv,
|
|
&mcpsdk.Tool{Name: "search_code", Description: "Semantic code search over the workspace's indexed files; returns ranked file:line snippets."},
|
|
func(ctx context.Context, _ *mcpsdk.CallToolRequest, in searchCodeIn) (*mcpsdk.CallToolResult, searchCodeOut, error) {
|
|
if err := CheckToolFromCtx(ctx, "search_code"); err != nil {
|
|
return nil, searchCodeOut{}, err
|
|
}
|
|
c, err := deps.Caller.Resolve(ctx)
|
|
if err != nil {
|
|
return nil, searchCodeOut{}, err
|
|
}
|
|
w, _, err := scopedWorkspace(ctx, deps, c, in.WorkspaceID)
|
|
if err != nil {
|
|
return nil, searchCodeOut{}, err
|
|
}
|
|
snippets, err := deps.CodeIndex.Search(ctx, w.ID, codeindex.SearchQuery{
|
|
Text: in.Query, TopK: in.TopK, PathPrefix: in.PathPrefix, Lang: in.Lang,
|
|
})
|
|
if err != nil {
|
|
return nil, searchCodeOut{}, err
|
|
}
|
|
out := searchCodeOut{Results: make([]codeSnippet, 0, len(snippets))}
|
|
for _, s := range snippets {
|
|
out.Results = append(out.Results, codeSnippet{
|
|
File: s.FilePath, StartLine: s.StartLine, EndLine: s.EndLine,
|
|
Snippet: s.Content, Score: s.Score, CommitSHA: s.CommitSHA,
|
|
})
|
|
}
|
|
return nil, out, nil
|
|
})
|
|
}
|
|
|
|
// ===== reindex_workspace =====
|
|
|
|
type reindexWorkspaceIn struct {
|
|
WorkspaceID string `json:"workspace_id" jsonschema:"workspace UUID"`
|
|
}
|
|
|
|
type reindexWorkspaceOut struct {
|
|
OK bool `json:"ok"`
|
|
RunID string `json:"run_id"`
|
|
}
|
|
|
|
func registerReindexWorkspace(srv *mcpsdk.Server, deps ServerDeps) {
|
|
mcpsdk.AddTool(srv,
|
|
&mcpsdk.Tool{Name: "reindex_workspace", Description: "Enqueue a fresh code index build at the workspace's current main HEAD."},
|
|
func(ctx context.Context, _ *mcpsdk.CallToolRequest, in reindexWorkspaceIn) (*mcpsdk.CallToolResult, reindexWorkspaceOut, error) {
|
|
if err := CheckToolFromCtx(ctx, "reindex_workspace"); err != nil {
|
|
return nil, reindexWorkspaceOut{}, err
|
|
}
|
|
c, err := deps.Caller.Resolve(ctx)
|
|
if err != nil {
|
|
return nil, reindexWorkspaceOut{}, err
|
|
}
|
|
w, _, err := scopedWorkspace(ctx, deps, c, in.WorkspaceID)
|
|
if err != nil {
|
|
return nil, reindexWorkspaceOut{}, err
|
|
}
|
|
run, err := deps.CodeIndex.EnqueueBuild(ctx, w.ID, nil, "")
|
|
if err != nil {
|
|
return nil, reindexWorkspaceOut{}, err
|
|
}
|
|
return nil, reindexWorkspaceOut{OK: true, RunID: run.ID.String()}, nil
|
|
})
|
|
}
|