diff --git a/config.example.toml b/config.example.toml index 67ffa9812..35b3680e9 100644 --- a/config.example.toml +++ b/config.example.toml @@ -109,6 +109,10 @@ GITHUB_PERSONAL_ACCESS_TOKEN = "" # Pass through from host # tools = ["*"] # Allow all tools (default) # tools = ["read_file", "list_files"] # Allow specific tools only +# Optional: Prompt filtering +# prompts = ["*"] # Allow all prompts (default) +# prompts = ["summarize", "translate"] # Allow specific prompts only + # Optional: Per-tool response filters (jq expressions) # Transforms tool response payloads before they are returned and before large-payload # preview/schema processing runs. diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index e666d6d5d..be8656983 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -230,6 +230,12 @@ Non-absolute entries and the filesystem root (`/`) are ignored. When the resulti - Use `["*"]` (wildcard) to allow all tools (default behavior when field is omitted) - Example: `["get_file_contents", "search_code"]` (only these tools are accessible) +- **`prompts`** (optional): List of prompt names to allow for this server + - Enforced at runtime: prompts not in this list are hidden from `prompts/list` responses and rejected if invoked via `prompts/get` + - Use `["*"]` (wildcard) to allow all prompts (default behavior when field is omitted) + - Prompt retrieval runs through session initialization, per-agent server authorization, and the configured DIFC guard pipeline + - Example: `["assign_coding_agent", "issue_to_fix_workflow"]` + - **`tool_response_filters`** (optional): Per-tool jq expressions that transform tool response data before it is returned to the agent and before large-payload preview/schema processing runs - Map key: tool name; map value: jq expression string - Expressions are validated at startup (compile check); invalid jq causes startup/config validation failure diff --git a/docs/GUARD_RESPONSE_LABELING.md b/docs/GUARD_RESPONSE_LABELING.md index e3b834555..fd883ca5d 100644 --- a/docs/GUARD_RESPONSE_LABELING.md +++ b/docs/GUARD_RESPONSE_LABELING.md @@ -1,6 +1,6 @@ # Guard Response Labeling -This document describes how guards label responses for DIFC (Decentralized Information Flow Control) enforcement in the MCP Gateway. +This document describes how guards label responses for DIFC (Decentralized Information Flow Control) enforcement in the MCP Gateway. The same reference-monitor phases mediate both `tools/call` and `prompts/get`; prompt names are passed to `LabelResource()` and `LabelResponse()` as the guarded operation name. When a prompt guard returns collection labels, the collection is interpreted as labels for the prompt's `messages` array and inaccessible messages are removed or cause the response to be blocked according to the active enforcement mode. ## DIFC Label Rules diff --git a/docs/RELEASE_NOTES.md b/docs/RELEASE_NOTES.md index 53e96da1b..41086e0bf 100644 --- a/docs/RELEASE_NOTES.md +++ b/docs/RELEASE_NOTES.md @@ -4,6 +4,7 @@ ### 🔒 Security +- **MCP prompt retrieval is now guarded**: `prompts/get` now requires session initialization and passes through per-agent server authorization, per-server prompt allowlisting, guard initialization, coarse DIFC checks, response labeling/filtering, and propagate-mode label accumulation. Prompt responses can no longer bypass the reference monitor through direct backend forwarding. - **DIFC nil-label responses fail closed**: Unified MCP calls now fall back to the coarse access decision when a guard returns no fine-grained response labels. A coarse-denied read is withheld instead of returning the original backend payload; coarse-allowed responses and propagate-mode reads retain their existing behavior. - **Custom guard compatibility**: Custom and third-party guards that intentionally return no response labels can no longer rely on the unified server passing through data after a coarse denial. Guards that need partial results for coarse-denied reads must provide fine-grained response labels; otherwise callers receive the standard DIFC access error. diff --git a/internal/config/config_core.go b/internal/config/config_core.go index e9163b859..88ed29504 100644 --- a/internal/config/config_core.go +++ b/internal/config/config_core.go @@ -345,6 +345,9 @@ type ServerConfig struct { // Tools is an optional list of tools to filter/expose Tools []string `toml:"tools" json:"tools,omitempty"` + // Prompts is an optional list of prompts to filter/expose. + Prompts []string `toml:"prompts" json:"prompts,omitempty"` + // ToolResponseFilters configures per-tool jq expressions that transform tool // response data before it is returned to the agent and before large-payload // preview/schema processing runs. diff --git a/internal/config/config_stdin.go b/internal/config/config_stdin.go index b63511be2..df5809528 100644 --- a/internal/config/config_stdin.go +++ b/internal/config/config_stdin.go @@ -151,6 +151,9 @@ type StdinServerConfig struct { // Tools is an optional list of tools to filter/expose Tools []string `json:"tools,omitempty"` + // Prompts is an optional list of prompts to filter/expose + Prompts []string `json:"prompts,omitempty"` + // ToolResponseFilters configures per-tool jq expressions that transform tool // response data before it is returned to the agent. ToolResponseFilters map[string]string `json:"tool_response_filters,omitempty"` @@ -245,6 +248,7 @@ func (s *StdinServerConfig) UnmarshalJSON(data []byte) error { "url": true, "headers": true, "tools": true, + "prompts": true, "tool_response_filters": true, "registry": true, "guard-policies": true, @@ -690,6 +694,7 @@ func buildStdioServerConfigWithRuntime(name string, server *StdinServerConfig, r // between HTTP and stdio server configurations. func applyCommonServerConfigFields(cfg *ServerConfig, src *StdinServerConfig) { cfg.Tools = src.Tools + cfg.Prompts = src.Prompts cfg.ToolResponseFilters = src.ToolResponseFilters cfg.Registry = src.Registry cfg.GuardPolicies = src.GuardPolicies diff --git a/internal/config/config_stdin_test.go b/internal/config/config_stdin_test.go index 51aeaa190..d40d0c76b 100644 --- a/internal/config/config_stdin_test.go +++ b/internal/config/config_stdin_test.go @@ -25,7 +25,8 @@ func TestConvertStdinServerConfig_StdioServer(t *testing.T) { "PASSTHROUGH": "", "ANOTHER_VAR": "abc", }, - Tools: []string{"tool1", "tool2"}, + Tools: []string{"tool1", "tool2"}, + Prompts: []string{"prompt1", "prompt2"}, } result, err := convertStdinServerConfig("test-server", server, nil) @@ -69,6 +70,7 @@ func TestConvertStdinServerConfig_StdioServer(t *testing.T) { // Check tools assert.Equal([]string{"tool1", "tool2"}, result.Tools) + assert.Equal([]string{"prompt1", "prompt2"}, result.Prompts) } // TestConvertStdinServerConfig_HTTPServer tests HTTP server conversion. @@ -83,7 +85,8 @@ func TestConvertStdinServerConfig_HTTPServer(t *testing.T) { "Authorization": "******", "X-Custom": "value", }, - Tools: []string{"search", "fetch"}, + Tools: []string{"search", "fetch"}, + Prompts: []string{"summarize"}, } result, err := convertStdinServerConfig("http-server", server, nil) @@ -95,6 +98,7 @@ func TestConvertStdinServerConfig_HTTPServer(t *testing.T) { assert.Equal("******", result.Headers["Authorization"]) assert.Equal("value", result.Headers["X-Custom"]) assert.Equal([]string{"search", "fetch"}, result.Tools) + assert.Equal([]string{"summarize"}, result.Prompts) // HTTP servers should not have Command or Args assert.Empty(result.Command) diff --git a/internal/server/prompt_call.go b/internal/server/prompt_call.go new file mode 100644 index 000000000..457c33c1a --- /dev/null +++ b/internal/server/prompt_call.go @@ -0,0 +1,173 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/github/gh-aw-mcpg/internal/difc" + "github.com/github/gh-aw-mcpg/internal/guard" + "github.com/github/gh-aw-mcpg/internal/logger" + "github.com/github/gh-aw-mcpg/internal/mcp" + "github.com/github/gh-aw-mcpg/internal/sanitize" + "github.com/github/gh-aw-mcpg/internal/util" + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func (us *UnifiedServer) isPromptAllowed(serverID, promptName string) bool { + set, ok := us.allowedPromptSets[serverID] + if !ok || set == nil { + return true + } + allowed := set[promptName] + if !allowed { + logUnified.Printf("isPromptAllowed: prompt blocked by allowlist: serverID=%s, promptName=%s", serverID, promptName) + } + return allowed +} + +// callBackendPrompt retrieves a backend prompt through the same DIFC reference-monitor +// phases used for tool responses. +func (us *UnifiedServer) callBackendPrompt(ctx context.Context, serverID, promptName string, args map[string]string) (*sdk.GetPromptResult, error) { + logUnified.Printf("callBackendPrompt: serverID=%s, promptName=%s", serverID, promptName) + + toolTimeout := 0 + if us.cfg != nil { + if serverCfg, ok := us.cfg.Servers[serverID]; ok && serverCfg != nil && serverCfg.ToolTimeout > 0 { + toolTimeout = serverCfg.ToolTimeout + } else if us.cfg.Gateway != nil && us.cfg.Gateway.ToolTimeout > 0 { + toolTimeout = us.cfg.Gateway.ToolTimeout + } + } + if toolTimeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(toolTimeout)*time.Second) + defer cancel() + } + + sessionID := us.getSessionID(ctx) + agentID := guard.GetAgentIDFromContext(ctx) + if us.agentPoliciesEnforced() && !us.agentCanAccessServer(agentID, serverID) { + return nil, fmt.Errorf("prompt %q on server %q is not permitted for this agent", promptName, serverID) + } + if !us.isPromptAllowed(serverID, promptName) { + return nil, fmt.Errorf("prompt %q is not in the allowed-prompts list for this server", promptName) + } + + g, err := us.guardForSession(ctx, sessionID, serverID) + if err != nil { + return nil, fmt.Errorf("failed to create isolated guard session: %w", err) + } + backendCaller := &guardBackendCaller{ + server: us, + serverID: serverID, + ctx: ctx, + } + enforcementMode, err := us.ensureGuardInitialized(ctx, sessionID, serverID, g, backendCaller) + if err != nil { + return nil, fmt.Errorf("guard session initialization failed: %w", err) + } + + requestEvaluator := difc.NewEvaluatorWithMode(enforcementMode) + pipelineIn := guard.PipelineInput{ + AgentID: agentID, + ToolName: promptName, + Args: args, + Guard: g, + Evaluator: requestEvaluator, + AgentRegistry: us.AgentRegistry, + Capabilities: us.Capabilities, + EnforcementMode: enforcementMode, + BackendCaller: backendCaller, + SensitiveLogging: sanitize.ShouldRedactPayload( + mcp.IsEnclaveSession(ctx), + ), + } + ctx, pre, err := guard.RunPipelinePrePhases(ctx, pipelineIn) + if err != nil { + if denied, detailedErr := guard.HandlePrePhaseError(err); denied != nil { + logCoarseDIFCDenial(ctx, serverID, promptName, denied) + return nil, detailedErr + } + return nil, fmt.Errorf("guard labeling failed: %w", err) + } + + backendResult, err := executeBackendRequest[interface{}](ctx, us.launcher, serverID, sessionID, "prompts/get", map[string]interface{}{ + "name": promptName, + "arguments": args, + }) + if err != nil { + return nil, err + } + + labeledData, err := guard.RunPipelinePhase4(ctx, pipelineIn, pre, backendResult) + if err != nil { + return nil, fmt.Errorf("response labeling failed: %w", err) + } + filterResult, err := difc.FilterAndConvertLabeledData( + requestEvaluator, + pre.AgentLabels.Secrecy, + pre.AgentLabels.Integrity, + pre.Operation, + labeledData, + enforcementMode, + ) + if err != nil { + return nil, fmt.Errorf("failed to convert labeled data: %w", err) + } + if filterResult.Blocked { + return nil, fmt.Errorf("DIFC policy violation: prompt response contains data that is not accessible to this agent") + } + + finalResult := backendResult + if labeledData != nil { + finalResult = filterResult.FinalResult + if _, ok := labeledData.(*difc.CollectionLabeledData); ok { + resultMap, mapOK := backendResult.(map[string]interface{}) + if !mapOK { + return nil, fmt.Errorf("prompt response labeling returned a collection for non-object response data") + } + messages, messagesOK := resultMap["messages"] + if !messagesOK { + return nil, fmt.Errorf("prompt response labeling returned a collection but the response has no messages") + } + if _, ok := messages.([]interface{}); !ok { + return nil, fmt.Errorf("prompt response messages have unexpected type %T", messages) + } + filteredMap := make(map[string]interface{}, len(resultMap)) + for key, value := range resultMap { + filteredMap[key] = value + } + filteredMap["messages"] = filterResult.FinalResult + finalResult = filteredMap + } + } else if !pre.EvalResult.IsAllowed() { + denied := &guard.PipelineAccessDenied{ + EvalResult: pre.EvalResult, + Resource: pre.Resource, + AgentLabels: pre.AgentLabels, + } + logCoarseDIFCDenial(ctx, serverID, promptName, denied) + return nil, difc.FormatViolationError( + pre.EvalResult, + pre.AgentLabels.Secrecy, + pre.AgentLabels.Integrity, + pre.Resource, + ) + } + + guard.RunPipelinePhase6(pre, labeledData, enforcementMode) + + resultJSON, err := json.Marshal(finalResult) + if err != nil { + return nil, fmt.Errorf("failed to marshal filtered prompt result: %w", err) + } + var result sdk.GetPromptResult + if err := json.Unmarshal(resultJSON, &result); err != nil { + return nil, fmt.Errorf("failed to parse filtered prompt result: %w", err) + } + logger.LogInfo("client", "MCP prompt response mediated, session=%s, prompt=%s", + util.FormatSessionIDForLog(sessionID), promptName) + return &result, nil +} diff --git a/internal/server/prompt_call_test.go b/internal/server/prompt_call_test.go new file mode 100644 index 000000000..cf5fa6bbe --- /dev/null +++ b/internal/server/prompt_call_test.go @@ -0,0 +1,333 @@ +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/difc" + "github.com/github/gh-aw-mcpg/internal/guard" + "github.com/github/gh-aw-mcpg/internal/launcher" + sdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type promptDIFCTestGuard struct { + mode difc.EnforcementMode + resource *difc.LabeledResource + operation difc.OperationType + response difc.LabeledData + labelAgent atomic.Int32 + labelResource atomic.Int32 + labelResponse atomic.Int32 +} + +func (g *promptDIFCTestGuard) Name() string { return "prompt-difc-test" } + +func (g *promptDIFCTestGuard) LabelAgent(context.Context, interface{}, guard.BackendCaller, *difc.Capabilities) (*guard.LabelAgentResult, error) { + g.labelAgent.Add(1) + return &guard.LabelAgentResult{ + Agent: guard.AgentLabelsPayload{}, + DIFCMode: g.mode.String(), + NormalizedPolicy: map[string]interface{}{ + "scope_kind": "all", + "min-integrity": "none", + }, + }, nil +} + +func (g *promptDIFCTestGuard) LabelResource(context.Context, string, interface{}, guard.BackendCaller, *difc.Capabilities) (*difc.LabeledResource, difc.OperationType, error) { + g.labelResource.Add(1) + if g.resource == nil { + return difc.NewLabeledResource("prompt"), difc.OperationRead, nil + } + return g.resource, g.operation, nil +} + +func (g *promptDIFCTestGuard) LabelResponse(context.Context, string, interface{}, guard.BackendCaller, *difc.Capabilities) (difc.LabeledData, error) { + g.labelResponse.Add(1) + return g.response, nil +} + +func newPromptDIFCBackend(t *testing.T, promptName string, promptResult map[string]interface{}, getCount *atomic.Int32) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + var req map[string]interface{} + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + switch req["method"] { + case "initialize": + w.Header().Set("Mcp-Session-Id", "prompt-difc-backend-session") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req["id"], + "result": map[string]interface{}{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]interface{}{ + "prompts": map[string]interface{}{}, + }, + "serverInfo": map[string]interface{}{"name": "prompt-difc-backend", "version": "1.0"}, + }, + }) + case "notifications/initialized": + w.WriteHeader(http.StatusAccepted) + case "tools/list": + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req["id"], + "result": map[string]interface{}{"tools": []interface{}{}}, + }) + case "prompts/list": + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req["id"], + "result": map[string]interface{}{ + "prompts": []map[string]interface{}{{"name": promptName}}, + }, + }) + case "prompts/get": + if getCount != nil { + getCount.Add(1) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req["id"], + "result": promptResult, + }) + default: + w.WriteHeader(http.StatusNotFound) + } + })) +} + +func newPromptDIFCTestServer(t *testing.T, backendURL string, prompts []string, g guard.Guard) (*UnifiedServer, *launcher.Launcher) { + t.Helper() + cfg := &config.Config{ + DIFCMode: difc.EnforcementStrict.String(), + Gateway: &config.GatewayConfig{ + ToolTimeout: 60, + PayloadDir: t.TempDir(), + }, + Servers: map[string]*config.ServerConfig{ + "prompts-server": { + Type: "http", + URL: backendURL, + Prompts: prompts, + }, + }, + GuardPolicy: &config.GuardPolicy{ + AllowOnly: &config.AllowOnlyPolicy{ + Repos: "public", + MinIntegrity: config.IntegrityNone, + }, + }, + GuardPolicySource: "test", + } + l := launcher.New(context.Background(), cfg) + t.Cleanup(func() { l.Close() }) + us := promptHandlerTestServer(t, "prompt-difc-test-server", cfg, l) + us.guardRegistry.Register("prompts-server", g) + return us, l +} + +func promptCallContext(sessionID string) context.Context { + ctx := context.WithValue(context.Background(), SessionIDContextKey, sessionID) + return guard.SetAgentIDInContext(ctx, sessionID) +} + +func TestPromptHandlerRunsSessionAndDIFCPipeline(t *testing.T) { + var getCount atomic.Int32 + backend := newPromptDIFCBackend(t, "greet", map[string]interface{}{ + "description": "greeting", + "messages": []map[string]interface{}{{ + "role": "user", + "content": map[string]interface{}{"type": "text", "text": "hello"}, + }}, + }, &getCount) + defer backend.Close() + + g := &promptDIFCTestGuard{mode: difc.EnforcementStrict} + us, l := newPromptDIFCTestServer(t, backend.URL, nil, g) + conn, err := launcher.GetOrLaunch(l, "prompts-server") + require.NoError(t, err) + require.NoError(t, us.registerPromptsFromBackend(context.Background(), "prompts-server", conn)) + + serverTransport, clientTransport := sdk.NewInMemoryTransports() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + go func() { _ = us.server.Run(ctx, serverTransport) }() + + client := sdk.NewClient(&sdk.Implementation{Name: "prompt-difc-client", Version: "1.0"}, nil) + clientSession, err := client.Connect(ctx, clientTransport, nil) + require.NoError(t, err) + defer clientSession.Close() + + result, err := clientSession.GetPrompt(ctx, &sdk.GetPromptParams{Name: "prompts-server___greet"}) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, int32(1), getCount.Load()) + assert.Equal(t, int32(1), g.labelAgent.Load()) + assert.Equal(t, int32(1), g.labelResource.Load()) + assert.Equal(t, int32(1), g.labelResponse.Load()) + us.sessionMu.RLock() + _, sessionCreated := us.sessions["default"] + us.sessionMu.RUnlock() + assert.True(t, sessionCreated, "prompt handler must initialize the caller session before dispatch") +} + +func TestCallBackendPromptRejectsPromptOutsideAllowlist(t *testing.T) { + var getCount atomic.Int32 + backend := newPromptDIFCBackend(t, "blocked", map[string]interface{}{ + "messages": []interface{}{}, + }, &getCount) + defer backend.Close() + + us, _ := newPromptDIFCTestServer(t, backend.URL, []string{"allowed"}, &promptDIFCTestGuard{mode: difc.EnforcementStrict}) + ctx := promptCallContext("allowlist-session") + require.NoError(t, us.requireSession(ctx)) + + result, err := us.callBackendPrompt(ctx, "prompts-server", "blocked", nil) + + require.Error(t, err) + assert.Nil(t, result) + require.ErrorContains(t, err, "allowed-prompts") + assert.Zero(t, getCount.Load(), "denied prompt must not reach the backend") +} + +func TestCallBackendPromptEnforcesAgentServerAuthorization(t *testing.T) { + var getCount atomic.Int32 + backend := newPromptDIFCBackend(t, "restricted", map[string]interface{}{ + "messages": []interface{}{}, + }, &getCount) + defer backend.Close() + + us, _ := newPromptDIFCTestServer(t, backend.URL, nil, &promptDIFCTestGuard{mode: difc.EnforcementStrict}) + us.cfg.Gateway.AgentIDs = []string{"denied-agent"} + us.cfg.Gateway.AgentPolicies = map[string]*config.AgentPolicy{ + "denied-agent": {Servers: []string{}}, + } + ctx := promptCallContext("denied-agent") + require.NoError(t, us.requireSession(ctx)) + + result, err := us.callBackendPrompt(ctx, "prompts-server", "restricted", nil) + + require.Error(t, err) + assert.Nil(t, result) + require.ErrorContains(t, err, "not permitted for this agent") + assert.Zero(t, getCount.Load(), "unauthorized prompt must not reach the backend") +} + +func TestCallBackendPromptCoarseDenyPreventsBackendExecution(t *testing.T) { + var getCount atomic.Int32 + backend := newPromptDIFCBackend(t, "write-like", map[string]interface{}{ + "messages": []interface{}{}, + }, &getCount) + defer backend.Close() + + resource := difc.NewLabeledResource("protected prompt") + resource.Integrity.Label.Add("approved") + g := &promptDIFCTestGuard{ + mode: difc.EnforcementStrict, + resource: resource, + operation: difc.OperationWrite, + } + us, _ := newPromptDIFCTestServer(t, backend.URL, nil, g) + ctx := promptCallContext("coarse-deny-session") + require.NoError(t, us.requireSession(ctx)) + + result, err := us.callBackendPrompt(ctx, "prompts-server", "write-like", nil) + + require.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "DIFC") + assert.Zero(t, getCount.Load(), "coarse-denied prompt must not reach the backend") + assert.Equal(t, int32(1), g.labelResource.Load()) + assert.Zero(t, g.labelResponse.Load()) +} + +func TestCallBackendPromptFiltersMessagesWithoutDisclosingDeniedContent(t *testing.T) { + var getCount atomic.Int32 + backend := newPromptDIFCBackend(t, "mixed", map[string]interface{}{ + "description": "mixed prompt", + "messages": []map[string]interface{}{ + {"role": "user", "content": map[string]interface{}{"type": "text", "text": "public message"}}, + {"role": "user", "content": map[string]interface{}{"type": "text", "text": "TOP SECRET"}}, + }, + }, &getCount) + defer backend.Close() + + publicLabels := difc.NewLabeledResource("public message") + privateLabels := difc.NewLabeledResource("private message") + privateLabels.Secrecy.Label.Add("private:restricted/repo") + g := &promptDIFCTestGuard{ + mode: difc.EnforcementFilter, + operation: difc.OperationRead, + response: &difc.CollectionLabeledData{Items: []difc.LabeledItem{ + {Data: map[string]interface{}{"role": "user", "content": map[string]interface{}{"type": "text", "text": "public message"}}, Labels: publicLabels}, + {Data: map[string]interface{}{"role": "user", "content": map[string]interface{}{"type": "text", "text": "TOP SECRET"}}, Labels: privateLabels}, + }}, + } + us, _ := newPromptDIFCTestServer(t, backend.URL, nil, g) + ctx := promptCallContext("filter-session") + require.NoError(t, us.requireSession(ctx)) + + result, err := us.callBackendPrompt(ctx, "prompts-server", "mixed", nil) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, result.Messages, 1) + text, ok := result.Messages[0].Content.(*sdk.TextContent) + require.True(t, ok) + assert.Equal(t, "public message", text.Text) + resultJSON, err := json.Marshal(result) + require.NoError(t, err) + assert.NotContains(t, string(resultJSON), "TOP SECRET") + assert.Equal(t, int32(1), getCount.Load()) +} + +func TestCallBackendPromptPropagateAccumulatesLabels(t *testing.T) { + backend := newPromptDIFCBackend(t, "private", map[string]interface{}{ + "messages": []map[string]interface{}{{ + "role": "user", + "content": map[string]interface{}{"type": "text", "text": "private data"}, + }}, + }, nil) + defer backend.Close() + + resource := difc.NewLabeledResource("private prompt") + resource.Secrecy.Label.Add("private:org/repo") + g := &promptDIFCTestGuard{ + mode: difc.EnforcementPropagate, + resource: resource, + operation: difc.OperationRead, + } + us, _ := newPromptDIFCTestServer(t, backend.URL, nil, g) + ctx := promptCallContext("propagate-session") + require.NoError(t, us.requireSession(ctx)) + + result, err := us.callBackendPrompt(ctx, "prompts-server", "private", nil) + + require.NoError(t, err) + require.NotNil(t, result) + labels := us.AgentRegistry.GetOrCreate("propagate-session") + assert.Contains(t, labels.GetSecrecyTags(), difc.Tag("private:org/repo")) +} diff --git a/internal/server/register_prompts_test.go b/internal/server/register_prompts_test.go index ae6bf5bb9..ab1c7ef55 100644 --- a/internal/server/register_prompts_test.go +++ b/internal/server/register_prompts_test.go @@ -10,6 +10,8 @@ import ( "time" "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/difc" + "github.com/github/gh-aw-mcpg/internal/guard" "github.com/github/gh-aw-mcpg/internal/launcher" "github.com/github/gh-aw-mcpg/internal/mcp" sdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -110,6 +112,24 @@ func minimalPromptsTestServer(t *testing.T) *UnifiedServer { } } +func promptHandlerTestServer(t *testing.T, name string, cfg *config.Config, l *launcher.Launcher) *UnifiedServer { + t.Helper() + components, err := difc.NewComponents(cfg.DIFCMode, difc.EnforcementStrict) + require.NoError(t, err) + us := &UnifiedServer{ + server: newSDKServer(name, logUnified), + launcher: l, + sessions: make(map[string]*Session), + payloadDir: t.TempDir(), + allowedPromptSets: buildAllowedPromptSets(cfg), + guardRegistry: guard.NewRegistry(), + DIFCComponents: components, + cfg: cfg, + } + us.guardRegistry.Register("prompts-server", guard.NewNoopGuard()) + return us +} + // TestRegisterPromptsFromBackend_NoCapability verifies that a backend without prompts // capability causes registerPromptsFromBackend to return immediately without error. // This exercises the !BackendHasPromptsCapability() early-return branch. @@ -220,9 +240,28 @@ func TestRegisterPromptsFromBackend_RegistersPrompts(t *testing.T) { conn := connectStreamableBackend(t, srv) us := minimalPromptsTestServer(t) + us.allowedPromptSets = map[string]map[string]bool{ + "prompts-server": {"summarize": true}, + } err := us.registerPromptsFromBackend(context.Background(), "prompts-server", conn) - assert.NoError(t, err, "should return nil when prompts are successfully registered") + require.NoError(t, err, "should return nil when prompts are successfully registered") + + serverTransport, clientTransport := sdk.NewInMemoryTransports() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + go func() { _ = us.server.Run(ctx, serverTransport) }() + + client := sdk.NewClient(&sdk.Implementation{Name: "prompt-list-test-client", Version: "1.0"}, nil) + clientSession, err := client.Connect(ctx, clientTransport, nil) + require.NoError(t, err) + defer clientSession.Close() + + listResult, err := clientSession.ListPrompts(ctx, &sdk.ListPromptsParams{}) + require.NoError(t, err) + require.Len(t, listResult.Prompts, 1) + assert.Equal(t, "prompts-server___summarize", listResult.Prompts[0].Name, + "prompts outside the configured allowlist must be hidden from prompts/list") } // TestRegisterPromptsFromBackend_JSONRPCErrorResponse verifies that a JSON-RPC-level @@ -364,10 +403,7 @@ func TestRegisterPromptsFromBackend_PromptHandlerInvocation(t *testing.T) { require.NoError(t, err) require.True(t, conn.BackendHasPromptsCapability()) - us := &UnifiedServer{ - server: newSDKServer("test-prompts-handler", logUnified), - launcher: l, - } + us := promptHandlerTestServer(t, "test-prompts-handler", cfg, l) err = us.registerPromptsFromBackend(context.Background(), "prompts-server", conn) require.NoError(t, err, "prompt registration should succeed") @@ -491,10 +527,7 @@ func TestRegisterPromptsFromBackend_PromptHandlerBackendError(t *testing.T) { require.NoError(t, err) require.True(t, conn.BackendHasPromptsCapability()) - us := &UnifiedServer{ - server: newSDKServer("test-prompts-handler-error", logUnified), - launcher: l, - } + us := promptHandlerTestServer(t, "test-prompts-handler-error", cfg, l) err = us.registerPromptsFromBackend(context.Background(), "prompts-server", conn) require.NoError(t, err, "prompt registration should succeed even though the prompt itself will fail on invocation") diff --git a/internal/server/tool_registry.go b/internal/server/tool_registry.go index a11a93062..d07092313 100644 --- a/internal/server/tool_registry.go +++ b/internal/server/tool_registry.go @@ -455,6 +455,21 @@ func (us *UnifiedServer) registerPromptsFromBackend(ctx context.Context, serverI return nil } + if allowedSet, ok := us.allowedPromptSets[serverID]; ok { + n := 0 + for _, prompt := range listResult.Prompts { + if allowedSet[prompt.Name] { + listResult.Prompts[n] = prompt + n++ + } + } + if n < len(listResult.Prompts) { + logger.LogInfo("backend", "[allowed-prompts] Filtered %d prompts from %s: keeping %d of %d", + len(listResult.Prompts)-n, serverID, n, len(listResult.Prompts)) + } + listResult.Prompts = listResult.Prompts[:n] + } + // Register each prompt with a prefixed name so front-end clients can // distinguish prompts from different backends. promptNames := make([]string, 0, len(listResult.Prompts)) @@ -474,15 +489,17 @@ func (us *UnifiedServer) registerPromptsFromBackend(ctx context.Context, serverI us.server.AddPrompt(sdkPrompt, func(ctx context.Context, req *sdk.GetPromptRequest) (*sdk.GetPromptResult, error) { sessionID := us.getSessionID(ctx) - params := map[string]interface{}{ - "name": promptNameCopy, - "arguments": req.Params.Arguments, + if err := us.requireSession(ctx); err != nil { + logger.LogError("client", "MCP prompt call failed: session not initialized, session=%s, prompt=%s", + util.FormatSessionIDForLog(sessionID), promptNameCopy) + return nil, err } - result, err := executeBackendRequest[sdk.GetPromptResult](ctx, us.launcher, serverIDCopy, sessionID, "prompts/get", params) + + result, err := us.callBackendPrompt(ctx, serverIDCopy, promptNameCopy, req.Params.Arguments) if err != nil { return nil, fmt.Errorf("failed to get prompt %s from backend %s: %w", promptNameCopy, serverIDCopy, err) } - return &result, nil + return result, nil }) logUnified.Printf("Registered prompt: %s___%s", serverID, promptNameCopy) diff --git a/internal/server/tool_registry_helpers.go b/internal/server/tool_registry_helpers.go index a65f8225c..ffed5e036 100644 --- a/internal/server/tool_registry_helpers.go +++ b/internal/server/tool_registry_helpers.go @@ -175,6 +175,32 @@ func buildAllowedToolSets(cfg *config.Config) map[string]map[string]bool { return sets } +// buildAllowedPromptSets builds a per-server prompt allowlist for O(1) lookup. +// Servers with no Prompts list, or with a "*" entry, permit every prompt. +func buildAllowedPromptSets(cfg *config.Config) map[string]map[string]bool { + sets := make(map[string]map[string]bool) + if cfg == nil { + return sets + } + for serverID, serverCfg := range cfg.Servers { + if len(serverCfg.Prompts) == 0 { + continue + } + if hasWildcard(serverCfg.Prompts) { + logger.LogInfo("backend", "[allowed-prompts] Wildcard \"*\" configured for %s: allowing all prompts", serverID) + continue + } + set := make(map[string]bool, len(serverCfg.Prompts)) + for _, prompt := range serverCfg.Prompts { + set[prompt] = true + } + sets[serverID] = set + logToolRegistryHelpers.Printf("Built allowed prompt set for server %s: %d prompt(s) permitted", serverID, len(set)) + } + logToolRegistryHelpers.Printf("Built allowed prompt sets: %d server(s) with prompt restrictions", len(sets)) + return sets +} + // hasWildcard reports whether the tools list contains a "*" entry. func hasWildcard(tools []string) bool { for _, t := range tools { diff --git a/internal/server/unified.go b/internal/server/unified.go index 59f43ec84..badcf70a0 100644 --- a/internal/server/unified.go +++ b/internal/server/unified.go @@ -103,6 +103,10 @@ type UnifiedServer struct { // means all tools are permitted for that server. allowedToolSets map[string]map[string]bool + // allowedPromptSets holds a pre-computed set of allowed prompt names per server ID. + // A missing or nil entry means all prompts are permitted for that server. + allowedPromptSets map[string]map[string]bool + // circuitBreakers holds a per-backend rate-limit circuit breaker keyed by server ID. circuitBreakers map[string]*circuitBreaker circuitBreakersMu sync.RWMutex @@ -188,6 +192,7 @@ func NewUnified(ctx context.Context, cfg *config.Config) (*UnifiedServer, error) payloadPathPrefix: payloadPathPrefix, payloadSizeThreshold: payloadSizeThreshold, allowedToolSets: buildAllowedToolSets(cfg), + allowedPromptSets: buildAllowedPromptSets(cfg), circuitBreakers: buildCircuitBreakers(cfg), // Initialize DIFC components