708 lines
25 KiB
Go
708 lines
25 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"fire-safety-ymd/internal/domain"
|
|
"fire-safety-ymd/internal/service"
|
|
)
|
|
|
|
const testMCPToken = "0123456789abcdef0123456789abcdef"
|
|
|
|
func TestNewMCPHandlerRejectsUnsafeOptions(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
options MCPOptions
|
|
}{
|
|
{name: "short token", options: MCPOptions{AuthToken: "short", MaxBodyBytes: 1, ToolTimeout: time.Second}},
|
|
{name: "header control token", options: MCPOptions{AuthToken: strings.Repeat("a", 31) + "\n", MaxBodyBytes: 1, ToolTimeout: time.Second}},
|
|
{name: "oversized body setting", options: MCPOptions{AuthToken: testMCPToken, MaxBodyBytes: maximumMCPBody + 1, ToolTimeout: time.Second}},
|
|
{name: "oversized timeout", options: MCPOptions{AuthToken: testMCPToken, MaxBodyBytes: 1, ToolTimeout: maximumToolTimeout + time.Second}},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
if _, err := NewMCPHandler(&fakeMCPSpatialService{}, tt.options); err == nil {
|
|
t.Fatal("NewMCPHandler() error = nil")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPTransportGuards(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
configure func(*http.Request)
|
|
wantStatus int
|
|
wantErrorCode string
|
|
}{
|
|
{
|
|
name: "missing auth",
|
|
body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`,
|
|
configure: func(request *http.Request) {
|
|
request.Header.Del("Authorization")
|
|
},
|
|
wantStatus: http.StatusUnauthorized,
|
|
wantErrorCode: "MCP_AUTH_INVALID",
|
|
},
|
|
{
|
|
name: "wrong auth",
|
|
body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`,
|
|
configure: func(request *http.Request) {
|
|
request.Header.Set("Authorization", "Bearer wrong")
|
|
},
|
|
wantStatus: http.StatusUnauthorized,
|
|
wantErrorCode: "MCP_AUTH_INVALID",
|
|
},
|
|
{
|
|
name: "browser origin",
|
|
body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`,
|
|
configure: func(request *http.Request) {
|
|
request.Header.Set("Origin", "https://untrusted.example")
|
|
},
|
|
wantStatus: http.StatusForbidden,
|
|
wantErrorCode: "MCP_ORIGIN_FORBIDDEN",
|
|
},
|
|
{
|
|
name: "wrong content type",
|
|
body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`,
|
|
configure: func(request *http.Request) {
|
|
request.Header.Set("Content-Type", "text/plain")
|
|
},
|
|
wantStatus: http.StatusUnsupportedMediaType,
|
|
wantErrorCode: "MCP_CONTENT_TYPE_INVALID",
|
|
},
|
|
{
|
|
name: "invalid JSON",
|
|
body: `{"jsonrpc":`,
|
|
wantStatus: http.StatusBadRequest,
|
|
wantErrorCode: "MCP_REQUEST_INVALID",
|
|
},
|
|
{
|
|
name: "batch rejected",
|
|
body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"}]`,
|
|
wantStatus: http.StatusBadRequest,
|
|
wantErrorCode: "MCP_REQUEST_INVALID",
|
|
},
|
|
{
|
|
name: "invalid id",
|
|
body: `{"jsonrpc":"2.0","id":{},"method":"tools/list"}`,
|
|
wantStatus: http.StatusBadRequest,
|
|
wantErrorCode: "MCP_REQUEST_INVALID",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
request := authenticatedMCPRequest(tt.body)
|
|
if tt.configure != nil {
|
|
tt.configure(request)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != tt.wantStatus {
|
|
t.Fatalf("status = %d, want %d; body=%s", response.Code, tt.wantStatus, response.Body.String())
|
|
}
|
|
if got := response.Header().Get("X-Request-ID"); got == "" {
|
|
t.Fatal("X-Request-ID header is missing")
|
|
}
|
|
if !strings.Contains(response.Body.String(), `"code":"`+tt.wantErrorCode+`"`) {
|
|
t.Fatalf("response = %s, want code %s", response.Body.String(), tt.wantErrorCode)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPRejectsOversizedBody(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 32, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusRequestEntityTooLarge || !strings.Contains(response.Body.String(), "MCP_REQUEST_BODY_TOO_LARGE") {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestMCPInitializeAndInitializedNotification(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
|
|
initialize := authenticatedMCPRequest(`{
|
|
"jsonrpc":"2.0",
|
|
"id":"init-1",
|
|
"method":"initialize",
|
|
"params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1"}}
|
|
}`)
|
|
initializeResponse := httptest.NewRecorder()
|
|
handler.ServeHTTP(initializeResponse, initialize)
|
|
if initializeResponse.Code != http.StatusOK {
|
|
t.Fatalf("initialize status=%d body=%s", initializeResponse.Code, initializeResponse.Body.String())
|
|
}
|
|
if !strings.Contains(initializeResponse.Body.String(), `"protocolVersion":"2025-06-18"`) ||
|
|
!strings.Contains(initializeResponse.Body.String(), `"tools":{"listChanged":false}`) {
|
|
t.Fatalf("unexpected initialize response: %s", initializeResponse.Body.String())
|
|
}
|
|
|
|
initialized := authenticatedMCPRequest(`{"jsonrpc":"2.0","method":"notifications/initialized"}`)
|
|
initializedResponse := httptest.NewRecorder()
|
|
handler.ServeHTTP(initializedResponse, initialized)
|
|
if initializedResponse.Code != http.StatusAccepted || initializedResponse.Body.Len() != 0 {
|
|
t.Fatalf("initialized status=%d body=%q", initializedResponse.Code, initializedResponse.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestMCPProvenProfileCompletesToolCallLifecycle(t *testing.T) {
|
|
spatial := &fakeMCPSpatialService{}
|
|
handler := newTestMCPHandler(t, spatial, 4096, time.Second)
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantStatus int
|
|
wantBody string
|
|
}{
|
|
{
|
|
name: "initialize without configurable protocol version",
|
|
body: `{"jsonrpc":"2.0","id":1,"method":"initialize"}`,
|
|
wantStatus: http.StatusOK,
|
|
wantBody: `"protocolVersion":"2025-06-18"`,
|
|
},
|
|
{
|
|
name: "initialized notification",
|
|
body: `{"jsonrpc":"2.0","method":"notifications/initialized"}`,
|
|
wantStatus: http.StatusAccepted,
|
|
},
|
|
{
|
|
name: "list tools",
|
|
body: `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`,
|
|
wantStatus: http.StatusOK,
|
|
wantBody: `"tools"`,
|
|
},
|
|
{
|
|
name: "call a read-only tool",
|
|
body: `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"fire_safety_search_place_candidates","arguments":{"place_name":"观水镇","limit":3}}}`,
|
|
wantStatus: http.StatusOK,
|
|
wantBody: `"structuredContent"`,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
request := authenticatedMCPRequest(tt.body)
|
|
request.Header.Set("MCP-Protocol-Version", "superagent-unconfigured-version")
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != tt.wantStatus || (tt.wantBody != "" && !strings.Contains(response.Body.String(), tt.wantBody)) {
|
|
t.Fatalf("status=%d body=%q", response.Code, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
if spatial.called != toolSearchPlaceCandidates {
|
|
t.Fatalf("called=%q, want %q", spatial.called, toolSearchPlaceCandidates)
|
|
}
|
|
}
|
|
|
|
func TestMCPInitializeAcceptsSuperAgentProtocolShapes(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
params string
|
|
protocolHeader string
|
|
}{
|
|
{name: "missing params"},
|
|
{name: "null params", params: "null"},
|
|
{name: "missing version", params: `{}`},
|
|
{name: "empty version", params: `{"protocolVersion":""}`},
|
|
{name: "blank version", params: `{"protocolVersion":" "}`},
|
|
{name: "non-string version", params: `{"protocolVersion":20250618}`},
|
|
{name: "invalid params shape", params: `[]`},
|
|
{name: "direct", params: `{"protocolVersion":"2025-06-18"}`},
|
|
{name: "older", params: `{"protocolVersion":"2025-03-26"}`},
|
|
{name: "newer and mismatched header", params: `{"protocolVersion":"2025-11-25"}`, protocolHeader: "2025-03-26"},
|
|
{name: "unknown", params: `{"protocolVersion":"superagent-private-version"}`, protocolHeader: "superagent-header-version"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
body := `{"jsonrpc":"2.0","id":1,"method":"initialize"`
|
|
if tt.params != "" {
|
|
body += `,"params":` + tt.params
|
|
}
|
|
request := authenticatedMCPRequest(body + `}`)
|
|
if tt.protocolHeader != "" {
|
|
request.Header.Set("MCP-Protocol-Version", tt.protocolHeader)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
var decoded rpcResponse
|
|
if err := json.Unmarshal(response.Body.Bytes(), &decoded); err != nil {
|
|
t.Fatalf("decode response: %v; body=%s", err, response.Body.String())
|
|
}
|
|
if response.Code != http.StatusOK ||
|
|
!strings.Contains(response.Body.String(), `"protocolVersion":"2025-06-18"`) ||
|
|
decoded.Error != nil {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPProtocolHeaderDoesNotBlockSuperAgentRequests(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
header string
|
|
}{
|
|
{name: "missing"},
|
|
{name: "server version", header: mcpProtocolVersion},
|
|
{name: "older", header: "2025-03-26"},
|
|
{name: "newer", header: "2025-11-25"},
|
|
{name: "unknown", header: "superagent-private-version"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":2,"method":"tools/list"}`)
|
|
if tt.header == "" {
|
|
request.Header.Del("MCP-Protocol-Version")
|
|
} else {
|
|
request.Header.Set("MCP-Protocol-Version", tt.header)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
var decoded rpcResponse
|
|
if err := json.Unmarshal(response.Body.Bytes(), &decoded); err != nil {
|
|
t.Fatalf("decode response: %v; body=%s", err, response.Body.String())
|
|
}
|
|
if response.Code != http.StatusOK || decoded.Error != nil || !strings.Contains(response.Body.String(), `"tools"`) {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPProtocolHeaderDoesNotBlockInitializedNotification(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
header string
|
|
}{
|
|
{name: "missing"},
|
|
{name: "server version", header: mcpProtocolVersion},
|
|
{name: "older", header: "2025-03-26"},
|
|
{name: "newer", header: "2025-11-25"},
|
|
{name: "unknown", header: "superagent-private-version"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","method":"notifications/initialized"}`)
|
|
if tt.header == "" {
|
|
request.Header.Del("MCP-Protocol-Version")
|
|
} else {
|
|
request.Header.Set("MCP-Protocol-Version", tt.header)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusAccepted || response.Body.Len() != 0 {
|
|
t.Fatalf("status=%d body=%q", response.Code, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPProtocolHeaderDoesNotBlockToolCalls(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
header string
|
|
}{
|
|
{name: "missing"},
|
|
{name: "server version", header: mcpProtocolVersion},
|
|
{name: "older", header: "2025-03-26"},
|
|
{name: "newer", header: "2025-11-25"},
|
|
{name: "unknown", header: "superagent-private-version"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
spatial := &fakeMCPSpatialService{}
|
|
handler := newTestMCPHandler(t, spatial, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"fire_safety_resolve_incident_context","arguments":{"longitude":121.7,"latitude":37.2}}}`)
|
|
if tt.header == "" {
|
|
request.Header.Del("MCP-Protocol-Version")
|
|
} else {
|
|
request.Header.Set("MCP-Protocol-Version", tt.header)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
var decoded rpcResponse
|
|
if err := json.Unmarshal(response.Body.Bytes(), &decoded); err != nil {
|
|
t.Fatalf("decode response: %v; body=%s", err, response.Body.String())
|
|
}
|
|
if response.Code != http.StatusOK || decoded.Error != nil || spatial.called != toolResolveIncidentContext {
|
|
t.Fatalf("status=%d called=%q body=%s", response.Code, spatial.called, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPInitializeLogsFiniteCompatibilityResultWithoutInput(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
params string
|
|
wantResult string
|
|
}{
|
|
{name: "direct", params: `{"protocolVersion":"2025-06-18"}`, wantResult: "direct_success"},
|
|
{name: "other", params: `{"protocolVersion":"untrusted-client-version"}`, wantResult: "compatibility_success"},
|
|
{name: "missing", wantResult: "compatibility_success"},
|
|
{name: "non-string", params: `{"protocolVersion":20250618,"extra":"sensitive-value"}`, wantResult: "compatibility_success"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var logs strings.Builder
|
|
handler, err := NewMCPHandler(&fakeMCPSpatialService{}, MCPOptions{
|
|
AuthToken: testMCPToken,
|
|
MaxBodyBytes: 4096,
|
|
ToolTimeout: time.Second,
|
|
Logger: log.New(&logs, "", 0),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMCPHandler() error = %v", err)
|
|
}
|
|
body := `{"jsonrpc":"2.0","id":1,"method":"initialize"`
|
|
if tt.params != "" {
|
|
body += `,"params":` + tt.params
|
|
}
|
|
request := authenticatedMCPRequest(body + `}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if !strings.Contains(logs.String(), "operation=initialize result="+tt.wantResult) {
|
|
t.Fatalf("logs=%q, want result %s", logs.String(), tt.wantResult)
|
|
}
|
|
for _, sensitive := range []string{"untrusted-client-version", "sensitive-value", testMCPToken} {
|
|
if strings.Contains(logs.String(), sensitive) {
|
|
t.Fatalf("logs contain request input %q: %q", sensitive, logs.String())
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPListsSevenBoundedReadOnlyTools(t *testing.T) {
|
|
handler := newTestMCPHandler(t, &fakeMCPSpatialService{}, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":2,"method":"tools/list"}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusOK {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
var decoded struct {
|
|
Result struct {
|
|
Tools []struct {
|
|
Name string `json:"name"`
|
|
InputSchema map[string]any `json:"inputSchema"`
|
|
Annotations map[string]any `json:"annotations"`
|
|
} `json:"tools"`
|
|
} `json:"result"`
|
|
}
|
|
if err := json.Unmarshal(response.Body.Bytes(), &decoded); err != nil {
|
|
t.Fatalf("decode response: %v", err)
|
|
}
|
|
if len(decoded.Result.Tools) != 7 {
|
|
t.Fatalf("tool count=%d, want 7", len(decoded.Result.Tools))
|
|
}
|
|
for _, tool := range decoded.Result.Tools {
|
|
if tool.Annotations["readOnlyHint"] != true || tool.Annotations["destructiveHint"] != false || tool.Annotations["openWorldHint"] != false {
|
|
t.Fatalf("tool %s annotations=%#v", tool.Name, tool.Annotations)
|
|
}
|
|
properties, ok := tool.InputSchema["properties"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("tool %s lacks properties schema: %#v", tool.Name, tool.InputSchema)
|
|
}
|
|
if tool.Name == toolSearchPlaceCandidates {
|
|
if properties["place_name"] == nil || properties["longitude"] != nil || properties["latitude"] != nil {
|
|
t.Fatalf("place search schema = %#v", tool.InputSchema)
|
|
}
|
|
} else if properties["longitude"] == nil || properties["latitude"] == nil {
|
|
t.Fatalf("tool %s lacks bounded coordinate schema: %#v", tool.Name, tool.InputSchema)
|
|
}
|
|
for _, forbidden := range []string{"user_id", "tenant_id", "role", "scope_mode", "allowed_towns", "sql"} {
|
|
if _, exists := properties[forbidden]; exists {
|
|
t.Fatalf("tool %s exposes forbidden authorization/query field %s", tool.Name, forbidden)
|
|
}
|
|
}
|
|
}
|
|
for _, sensitiveColumn := range []string{"bpld", "bplddh", "csjdh", "fhzddclxfs", "zbrylxfs"} {
|
|
if strings.Contains(response.Body.String(), sensitiveColumn) {
|
|
t.Fatalf("tools/list exposes sensitive source column %s", sensitiveColumn)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestMCPDispatchesEverySpatialTool(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
arguments string
|
|
}{
|
|
{name: toolSearchPlaceCandidates, arguments: `{"place_name":"观水镇","limit":2}`},
|
|
{name: toolResolveIncidentContext, arguments: `{"longitude":121.7,"latitude":37.2}`},
|
|
{name: toolFindNearbyWaterSources, arguments: `{"longitude":121.7,"latitude":37.2,"radius_meters":1000,"limit":2}`},
|
|
{name: toolFindCommandPostCandidates, arguments: `{"longitude":121.7,"latitude":37.2}`},
|
|
{name: toolListNearbyAccessLines, arguments: `{"longitude":121.7,"latitude":37.2}`},
|
|
{name: toolGetResponsibleUnits, arguments: `{"longitude":121.7,"latitude":37.2}`},
|
|
{name: toolFindNearbyRiskAreas, arguments: `{"longitude":121.7,"latitude":37.2}`},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
spatial := &fakeMCPSpatialService{}
|
|
handler := newTestMCPHandler(t, spatial, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"` + tt.name + `","arguments":` + tt.arguments + `}}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusOK || spatial.called != tt.name {
|
|
t.Fatalf("status=%d called=%q body=%s", response.Code, spatial.called, response.Body.String())
|
|
}
|
|
var decoded struct {
|
|
Result struct {
|
|
Content []struct {
|
|
Text string `json:"text"`
|
|
} `json:"content"`
|
|
StructuredContent map[string]any `json:"structuredContent"`
|
|
IsError bool `json:"isError"`
|
|
} `json:"result"`
|
|
}
|
|
if err := json.Unmarshal(response.Body.Bytes(), &decoded); err != nil {
|
|
t.Fatalf("decode response: %v", err)
|
|
}
|
|
if decoded.Result.IsError || len(decoded.Result.Content) != 1 {
|
|
t.Fatalf("unexpected tool result: %s", response.Body.String())
|
|
}
|
|
var textContent map[string]any
|
|
if err := json.Unmarshal([]byte(decoded.Result.Content[0].Text), &textContent); err != nil {
|
|
t.Fatalf("text content is not mirrored JSON: %v", err)
|
|
}
|
|
if !reflect.DeepEqual(textContent, decoded.Result.StructuredContent) {
|
|
t.Fatalf("text content and structuredContent differ: text=%#v structured=%#v", textContent, decoded.Result.StructuredContent)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPToolErrorsAreStableAndSanitized(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
spatial *fakeMCPSpatialService
|
|
tool string
|
|
arguments string
|
|
timeout time.Duration
|
|
wantCode string
|
|
forbiddenText string
|
|
}{
|
|
{
|
|
name: "missing place name",
|
|
spatial: &fakeMCPSpatialService{},
|
|
tool: toolSearchPlaceCandidates,
|
|
arguments: `{"limit":2}`,
|
|
timeout: time.Second,
|
|
wantCode: "INVALID_ARGUMENT",
|
|
},
|
|
{
|
|
name: "unknown argument",
|
|
spatial: &fakeMCPSpatialService{},
|
|
tool: toolResolveIncidentContext,
|
|
arguments: `{"longitude":121,"latitude":37,"user_id":"admin"}`,
|
|
timeout: time.Second,
|
|
wantCode: "INVALID_ARGUMENT",
|
|
},
|
|
{
|
|
name: "unknown tool",
|
|
spatial: &fakeMCPSpatialService{},
|
|
tool: "execute_sql",
|
|
arguments: `{}`,
|
|
timeout: time.Second,
|
|
wantCode: "TOOL_NOT_FOUND",
|
|
},
|
|
{
|
|
name: "repository failure",
|
|
spatial: &fakeMCPSpatialService{failure: errors.New("postgres://secret-user:secret-password@internal/db")},
|
|
tool: toolFindNearbyWaterSources,
|
|
arguments: `{"longitude":121,"latitude":37}`,
|
|
timeout: time.Second,
|
|
wantCode: "DATA_SOURCE_UNAVAILABLE",
|
|
forbiddenText: "secret-password",
|
|
},
|
|
{
|
|
name: "tool timeout",
|
|
spatial: &fakeMCPSpatialService{waitForCancellation: true},
|
|
tool: toolFindNearbyWaterSources,
|
|
arguments: `{"longitude":121,"latitude":37}`,
|
|
timeout: 5 * time.Millisecond,
|
|
wantCode: "QUERY_TIMEOUT",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
handler := newTestMCPHandler(t, tt.spatial, 4096, tt.timeout)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"` + tt.tool + `","arguments":` + tt.arguments + `}}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusOK || !strings.Contains(response.Body.String(), `"isError":true`) || !strings.Contains(response.Body.String(), `"code":"`+tt.wantCode+`"`) {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
if tt.forbiddenText != "" && strings.Contains(response.Body.String(), tt.forbiddenText) {
|
|
t.Fatalf("response leaked internal error: %s", response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMCPDoesNotExecuteIDLessToolNotification(t *testing.T) {
|
|
spatial := &fakeMCPSpatialService{}
|
|
handler := newTestMCPHandler(t, spatial, 4096, time.Second)
|
|
request := authenticatedMCPRequest(`{"jsonrpc":"2.0","method":"tools/call","params":{"name":"fire_safety_resolve_incident_context","arguments":{"longitude":121,"latitude":37}}}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusAccepted || response.Body.Len() != 0 || spatial.called != "" {
|
|
t.Fatalf("status=%d called=%q body=%q", response.Code, spatial.called, response.Body.String())
|
|
}
|
|
}
|
|
|
|
func newTestMCPHandler(t *testing.T, spatial MCPSpatialService, maxBodyBytes int64, timeout time.Duration) *MCPHandler {
|
|
t.Helper()
|
|
handler, err := NewMCPHandler(spatial, MCPOptions{
|
|
AuthToken: testMCPToken,
|
|
MaxBodyBytes: maxBodyBytes,
|
|
ToolTimeout: timeout,
|
|
Logger: log.New(io.Discard, "", 0),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMCPHandler() error = %v", err)
|
|
}
|
|
return handler
|
|
}
|
|
|
|
func authenticatedMCPRequest(body string) *http.Request {
|
|
request := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
|
|
request.Header.Set("Authorization", "Bearer "+testMCPToken)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
request.Header.Set("MCP-Protocol-Version", mcpProtocolVersion)
|
|
return request
|
|
}
|
|
|
|
type fakeMCPSpatialService struct {
|
|
called string
|
|
failure error
|
|
waitForCancellation bool
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) prepare(ctx context.Context, tool string) error {
|
|
f.called = tool
|
|
if f.waitForCancellation {
|
|
<-ctx.Done()
|
|
return ctx.Err()
|
|
}
|
|
if f.failure != nil {
|
|
return errors.Join(service.ErrDataSourceUnavailable, f.failure)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) SearchPlaceCandidates(ctx context.Context, _ string, _ int) (service.QueryResult[[]domain.PlaceCandidate], error) {
|
|
if err := f.prepare(ctx, toolSearchPlaceCandidates); err != nil {
|
|
return service.QueryResult[[]domain.PlaceCandidate]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.PlaceCandidate{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) ResolveIncidentContext(ctx context.Context, _ domain.Coordinate) (service.QueryResult[[]domain.IncidentContext], error) {
|
|
if err := f.prepare(ctx, toolResolveIncidentContext); err != nil {
|
|
return service.QueryResult[[]domain.IncidentContext]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.IncidentContext{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) FindNearbyWaterSources(ctx context.Context, _ domain.Coordinate, _ float64, _ int) (service.QueryResult[[]domain.WaterSource], error) {
|
|
if err := f.prepare(ctx, toolFindNearbyWaterSources); err != nil {
|
|
return service.QueryResult[[]domain.WaterSource]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.WaterSource{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) FindCommandPostCandidates(ctx context.Context, _ domain.Coordinate, _ float64, _ int) (service.QueryResult[[]domain.CommandPostCandidate], error) {
|
|
if err := f.prepare(ctx, toolFindCommandPostCandidates); err != nil {
|
|
return service.QueryResult[[]domain.CommandPostCandidate]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.CommandPostCandidate{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) ListNearbyAccessLines(ctx context.Context, _ domain.Coordinate, _ float64, _ int) (service.QueryResult[[]domain.AccessLine], error) {
|
|
if err := f.prepare(ctx, toolListNearbyAccessLines); err != nil {
|
|
return service.QueryResult[[]domain.AccessLine]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.AccessLine{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) GetResponsibleUnits(ctx context.Context, _ domain.Coordinate) (service.QueryResult[[]domain.ResponsibleUnit], error) {
|
|
if err := f.prepare(ctx, toolGetResponsibleUnits); err != nil {
|
|
return service.QueryResult[[]domain.ResponsibleUnit]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.ResponsibleUnit{}), nil
|
|
}
|
|
|
|
func (f *fakeMCPSpatialService) FindNearbyRiskAreas(ctx context.Context, _ domain.Coordinate, _ float64, _ int) (service.QueryResult[[]domain.RiskArea], error) {
|
|
if err := f.prepare(ctx, toolFindNearbyRiskAreas); err != nil {
|
|
return service.QueryResult[[]domain.RiskArea]{}, err
|
|
}
|
|
return emptyMCPResult([]domain.RiskArea{}), nil
|
|
}
|
|
|
|
func emptyMCPResult[T any](data []T) service.QueryResult[[]T] {
|
|
return service.QueryResult[[]T]{
|
|
Status: "no_results",
|
|
Data: data,
|
|
Metadata: service.ResultMetadata{
|
|
GeneratedAt: "2026-09-04T00:00:00Z",
|
|
DataSources: []string{},
|
|
SpatialReference: "EPSG:4326",
|
|
ResultCount: 0,
|
|
},
|
|
Warnings: []string{},
|
|
}
|
|
}
|