Skip to content

Commit 7a8a370

Browse files
prachi-shah-harnessHarness
authored andcommitted
fix: [ML-1268]: prompt tools get and list (#151)
* remove redundant conditions * remove scope from list prompts * rename promptTools to prompt_tools.go * fix lint error * fix lint error * fix imports * fix lint error * fix imports * add register prompts to setup * add missing imports * modify unit test * merge master * rerun e2e * add test for prompt tools * fix: [ML-1268]: prompt tools get and list
1 parent da8a1f0 commit 7a8a370

File tree

5 files changed

+448
-1
lines changed

5 files changed

+448
-1
lines changed

pkg/harness/tools.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,9 @@ func initLegacyToolsets(config *config.Config, tsg *toolsets.ToolsetGroup) error
277277
if err := modules.RegisterSecrets(config, tsg); err != nil {
278278
return err
279279
}
280+
if err := modules.RegisterPromptTools(config, tsg); err != nil {
281+
return err
282+
}
280283
} else {
281284
// Register specified toolsets
282285
for _, toolset := range config.Toolsets {
@@ -381,6 +384,10 @@ func initLegacyToolsets(config *config.Config, tsg *toolsets.ToolsetGroup) error
381384
if err := modules.RegisterSecrets(config, tsg); err != nil {
382385
return err
383386
}
387+
case "prompts":
388+
if err := modules.RegisterPromptTools(config, tsg); err != nil {
389+
return err
390+
}
384391
}
385392
}
386393
}

pkg/harness/tools/prompt_tools.go

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
package tools
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"strings"
8+
9+
"github.com/harness/harness-mcp/cmd/harness-mcp-server/config"
10+
"github.com/mark3labs/mcp-go/mcp"
11+
"github.com/mark3labs/mcp-go/server"
12+
)
13+
14+
// ListPromptsTool creates a tool for listing prompts from the MCP server
15+
func ListPromptsTool(config *config.Config) (tool mcp.Tool, handler server.ToolHandlerFunc) {
16+
return mcp.NewTool("list_prompts",
17+
mcp.WithDescription("Lists available prompts from the MCP server"),
18+
mcp.WithString("prefix",
19+
mcp.Description("Optional prefix to filter prompts by name"),
20+
),
21+
),
22+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
23+
// Get MCP server from context
24+
mcpServer := server.ServerFromContext(ctx)
25+
if mcpServer == nil {
26+
return nil, fmt.Errorf("MCP server not found in context")
27+
}
28+
29+
// Extract prefix parameter if present
30+
prefix, err := OptionalParam[string](request, "prefix")
31+
if err != nil {
32+
return mcp.NewToolResultError(err.Error()), nil
33+
}
34+
35+
// Create a list prompts request
36+
listRequest := map[string]interface{}{
37+
"jsonrpc": "2.0",
38+
"id": "internal-tool-call",
39+
"method": "prompts/list",
40+
"params": map[string]interface{}{
41+
"cursor": "",
42+
},
43+
}
44+
45+
// Convert request to JSON
46+
listRequestBytes, err := json.Marshal(listRequest)
47+
if err != nil {
48+
return nil, fmt.Errorf("failed to marshal list request: %w", err)
49+
}
50+
51+
// Send request through HandleMessage
52+
listResponse := mcpServer.HandleMessage(ctx, listRequestBytes)
53+
54+
// Check for error response
55+
if errResp, isErr := listResponse.(mcp.JSONRPCError); isErr {
56+
return nil, fmt.Errorf("error listing prompts: %s", errResp.Error.Message)
57+
}
58+
59+
// Parse response to get prompts list
60+
jsonResp, ok := listResponse.(mcp.JSONRPCResponse)
61+
if !ok {
62+
return nil, fmt.Errorf("unexpected response type from list prompts")
63+
}
64+
65+
// Just return the raw result if we don't need filtering
66+
if prefix == "" {
67+
r, err := json.Marshal(jsonResp.Result)
68+
if err != nil {
69+
return nil, fmt.Errorf("failed to marshal result: %w", err)
70+
}
71+
return mcp.NewToolResultText(string(r)), nil
72+
}
73+
74+
// If we have a prefix, we need to extract and filter the prompts
75+
var listResult map[string]interface{}
76+
resultBytes, err := json.Marshal(jsonResp.Result)
77+
if err != nil {
78+
return nil, fmt.Errorf("failed to marshal result: %w", err)
79+
}
80+
81+
if err := json.Unmarshal(resultBytes, &listResult); err != nil {
82+
return nil, fmt.Errorf("failed to unmarshal result: %w", err)
83+
}
84+
85+
// Filter prompts by prefix if specified
86+
if promptsList, ok := listResult["prompts"].([]interface{}); ok {
87+
filtered := make([]interface{}, 0)
88+
for _, p := range promptsList {
89+
if prompt, ok := p.(map[string]interface{}); ok {
90+
if name, ok := prompt["name"].(string); ok {
91+
if strings.HasPrefix(name, prefix) {
92+
filtered = append(filtered, prompt)
93+
}
94+
}
95+
}
96+
}
97+
listResult["prompts"] = filtered
98+
}
99+
100+
r, err := json.Marshal(listResult)
101+
if err != nil {
102+
return nil, fmt.Errorf("failed to marshal filtered results: %w", err)
103+
}
104+
105+
return mcp.NewToolResultText(string(r)), nil
106+
}
107+
}
108+
109+
// GetPromptTool creates a tool for retrieving a single prompt from the MCP server
110+
func GetPromptTool(config *config.Config) (tool mcp.Tool, handler server.ToolHandlerFunc) {
111+
return mcp.NewTool("get_prompt",
112+
mcp.WithDescription("Retrieves a specific prompt from the MCP server by name"),
113+
mcp.WithString("prompt_name",
114+
mcp.Description("The name of the prompt to retrieve"),
115+
mcp.Required(),
116+
),
117+
mcp.WithString("mode",
118+
mcp.Description("Optional mode to retrieve a specific version of the prompt"),
119+
),
120+
),
121+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
122+
// Get MCP server from context
123+
mcpServer := server.ServerFromContext(ctx)
124+
if mcpServer == nil {
125+
return nil, fmt.Errorf("MCP server not found in context")
126+
}
127+
128+
// Extract prompt name parameter
129+
promptName, err := RequiredParam[string](request, "prompt_name")
130+
if err != nil {
131+
return mcp.NewToolResultError(err.Error()), nil
132+
}
133+
134+
// Extract optional mode parameter
135+
mode, err := OptionalParam[string](request, "mode")
136+
if err != nil {
137+
return mcp.NewToolResultError(err.Error()), nil
138+
}
139+
140+
// Create arguments map for the request
141+
arguments := make(map[string]interface{})
142+
if mode != "" {
143+
arguments["mode"] = mode
144+
}
145+
146+
// Create a get prompt request
147+
getRequest := map[string]interface{}{
148+
"jsonrpc": "2.0",
149+
"id": "internal-get-prompt",
150+
"method": "prompts/get",
151+
"params": map[string]interface{}{
152+
"name": promptName,
153+
"arguments": arguments,
154+
},
155+
}
156+
157+
// Convert to JSON
158+
getRequestBytes, err := json.Marshal(getRequest)
159+
if err != nil {
160+
return nil, fmt.Errorf("failed to marshal get prompt request: %w", err)
161+
}
162+
163+
// Send through HandleMessage
164+
getResponse := mcpServer.HandleMessage(ctx, getRequestBytes)
165+
166+
// Check for error
167+
if errResp, isErr := getResponse.(mcp.JSONRPCError); isErr {
168+
return nil, fmt.Errorf("error getting prompt: %s", errResp.Error.Message)
169+
}
170+
171+
// Parse response
172+
jsonResp, ok := getResponse.(mcp.JSONRPCResponse)
173+
if !ok {
174+
return nil, fmt.Errorf("unexpected response type from get prompt")
175+
}
176+
177+
// Return the result as JSON string
178+
r, err := json.Marshal(jsonResp.Result)
179+
if err != nil {
180+
return nil, fmt.Errorf("failed to marshal result: %w", err)
181+
}
182+
183+
return mcp.NewToolResultText(string(r)), nil
184+
}
185+
}

pkg/modules/core.go

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ func (m *CoreModule) Toolsets() []string {
5757
"chatbot",
5858
"settings",
5959
"secrets",
60+
"prompts",
6061
}
6162
}
6263

@@ -125,6 +126,11 @@ func (m *CoreModule) RegisterToolsets() error {
125126
if err != nil {
126127
return err
127128
}
129+
case "prompts":
130+
err := RegisterPromptTools(m.config, m.tsg)
131+
if err != nil {
132+
return err
133+
}
128134
}
129135
}
130136
return nil
@@ -429,7 +435,6 @@ func RegisterChatbot(config *config.Config, tsg *toolsets.ToolsetGroup) error {
429435
return nil
430436
}
431437

432-
// RegisterSettings registers the settings toolset
433438
func RegisterSettings(config *config.Config, tsg *toolsets.ToolsetGroup) error {
434439
// Determine the base URL and secret for settings
435440
baseURL := utils.BuildServiceURL(config, config.NgManagerBaseURL, config.BaseURL, "ng/api")
@@ -478,3 +483,16 @@ func RegisterSecrets(config *config.Config, tsg *toolsets.ToolsetGroup) error {
478483
tsg.AddToolset(secrets)
479484
return nil
480485
}
486+
487+
func RegisterPromptTools(config *config.Config, tsg *toolsets.ToolsetGroup) error {
488+
// Create the prompt toolset with both tools
489+
prompt := toolsets.NewToolset("prompt", "Harness MCP Prompts tools").
490+
AddReadTools(
491+
toolsets.NewServerTool(tools.GetPromptTool(config)),
492+
toolsets.NewServerTool(tools.ListPromptsTool(config)),
493+
)
494+
495+
// Add toolset to the group
496+
tsg.AddToolset(prompt)
497+
return nil
498+
}

0 commit comments

Comments
 (0)