diff --git a/internal/apiserver/authorized_surface.go b/internal/apiserver/authorized_surface.go index 64a8018f4..101eb74ad 100644 --- a/internal/apiserver/authorized_surface.go +++ b/internal/apiserver/authorized_surface.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "sort" + "strings" ) // ToolDescriptor is the internal, protocol-neutral form of a host tool. @@ -60,6 +61,9 @@ func validateToolDescriptor(tool ToolDescriptor) error { if input == nil || input["type"] != "object" { return fmt.Errorf("tool %q: input schema must have type object", tool.Name) } + if err := validateHeaderAnnotations(input); err != nil { + return fmt.Errorf("tool %q: invalid parameter header annotations: %w", tool.Name, err) + } if tool.OutputSchema != nil { var output any if err := json.Unmarshal(tool.OutputSchema, &output); err != nil { @@ -69,6 +73,51 @@ func validateToolDescriptor(tool ToolDescriptor) error { return nil } +func validateHeaderAnnotations(schema map[string]any) error { + seen := make(map[string]bool) + var walk func(map[string]any, string) error + walk = func(node map[string]any, prefix string) error { + properties, _ := node["properties"].(map[string]any) + for name, value := range properties { + property, _ := value.(map[string]any) + path := name + if prefix != "" { + path = prefix + "." + name + } + if annotation, ok := property["x-mcp-header"]; ok { + typeName, _ := property["type"].(string) + if typeName != "string" && typeName != "integer" && typeName != "boolean" { + return fmt.Errorf("property %q: x-mcp-header requires a primitive type", path) + } + header, ok := annotation.(string) + if !ok || header == "" || !validHTTPFieldName(header) { + return fmt.Errorf("property %q: invalid x-mcp-header value", path) + } + key := strings.ToLower(header) + if seen[key] { + return fmt.Errorf("property %q: duplicate x-mcp-header value %q", path, header) + } + seen[key] = true + } + if err := walk(property, path); err != nil { + return err + } + } + return nil + } + return walk(schema, "") +} + +func validHTTPFieldName(name string) bool { + for _, c := range name { + if !((c >= '0' && c <= '9') || (c >= 'A' && c <= 'Z') || + (c >= 'a' && c <= 'z') || strings.ContainsRune("!#$%&'*+-.^_`|~", c)) { + return false + } + } + return name != "" +} + func cloneToolDescriptor(tool ToolDescriptor) ToolDescriptor { tool.InputSchema = append(json.RawMessage(nil), tool.InputSchema...) tool.OutputSchema = append(json.RawMessage(nil), tool.OutputSchema...) diff --git a/internal/apiserver/embedded_mcp.go b/internal/apiserver/embedded_mcp.go new file mode 100644 index 000000000..9512798db --- /dev/null +++ b/internal/apiserver/embedded_mcp.go @@ -0,0 +1,232 @@ +package apiserver + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "sync" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// EmbeddedMCPConfig supplies the request-scoped host operations used by the +// public serveapi facade without introducing an internal-to-public import. +type EmbeddedMCPConfig struct { + AgentName string + AgentVersion string + Runner *CallbackRunner + Resolve func(context.Context, *http.Request) (any, error) + Tools func(context.Context, any) ([]ToolDescriptor, error) + Invoke func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) +} + +type embeddedMCPHandler struct { + config EmbeddedMCPConfig + transport http.Handler +} + +type embeddedMCPContext struct { + surface *AuthorizedSurface + session any + failure *embeddedCallbackFailure +} + +type embeddedCallbackFailure struct { + mu sync.Mutex + message string +} + +type embeddedMCPContextKey struct{} + +// NewEmbeddedMCPHandler builds the stateless SDK transport once. The server +// returned to it is still new for every authorized POST. +func NewEmbeddedMCPHandler(config EmbeddedMCPConfig) http.Handler { + h := &embeddedMCPHandler{config: config} + h.transport = mcp.NewStreamableHTTPHandler(h.serverForRequest, + &mcp.StreamableHTTPOptions{Stateless: true}) + return h +} + +func (h *embeddedMCPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + h.transport.ServeHTTP(w, r) + return + } + + body, err := io.ReadAll(io.LimitReader(r.Body, mcp.DefaultMaxRequestBodyBytes+1)) + if err != nil || len(body) > mcp.DefaultMaxRequestBodyBytes { + writeMCPEnvelope(w, requestID(body), "invalid MCP request body") + return + } + r.Body = io.NopCloser(bytes.NewReader(body)) + id := requestID(body) + + session, err := RunCallback(r.Context(), h.config.Runner, func(ctx context.Context) (any, error) { + return h.config.Resolve(ctx, r) + }) + if err != nil { + if status := authorizationStatus(err); status != 0 { + http.Error(w, err.Error(), status) + return + } + writeMCPCallbackError(w, id, err) + return + } + + descriptors, err := RunCallback(r.Context(), h.config.Runner, func(ctx context.Context) ([]ToolDescriptor, error) { + return h.config.Tools(ctx, session) + }) + if err != nil { + writeMCPCallbackError(w, id, err) + return + } + surface, err := callerSurface(descriptors) + if err != nil { + writeMCPEnvelope(w, id, err.Error()) + return + } + + failure := &embeddedCallbackFailure{} + ctx := context.WithValue(r.Context(), embeddedMCPContextKey{}, embeddedMCPContext{surface, session, failure}) + r = r.WithContext(ctx) + r.Body = io.NopCloser(bytes.NewReader(body)) + h.serveTransport(w, r, id) +} + +func (h *embeddedMCPHandler) serveTransport(w http.ResponseWriter, r *http.Request, id json.RawMessage) { + buffer := newBufferedResponseWriter() + defer func() { + if recover() != nil { + writeMCPEnvelope(w, id, "host tool registration failed") + } + }() + h.transport.ServeHTTP(buffer, r) + requestContext := r.Context().Value(embeddedMCPContextKey{}).(embeddedMCPContext) + requestContext.failure.mu.Lock() + message := requestContext.failure.message + requestContext.failure.mu.Unlock() + if message != "" { + writeMCPEnvelope(w, id, message) + return + } + for name, values := range buffer.header { + w.Header()[name] = append([]string(nil), values...) + } + w.WriteHeader(buffer.status) + _, _ = w.Write(buffer.body.Bytes()) +} + +type bufferedResponseWriter struct { + header http.Header + body bytes.Buffer + status int +} + +func newBufferedResponseWriter() *bufferedResponseWriter { + return &bufferedResponseWriter{header: make(http.Header), status: http.StatusOK} +} + +func (w *bufferedResponseWriter) Header() http.Header { return w.header } +func (w *bufferedResponseWriter) WriteHeader(status int) { w.status = status } +func (w *bufferedResponseWriter) Write(data []byte) (int, error) { return w.body.Write(data) } +func (w *bufferedResponseWriter) Flush() {} + +func (h *embeddedMCPHandler) serverForRequest(r *http.Request) *mcp.Server { + requestContext, ok := r.Context().Value(embeddedMCPContextKey{}).(embeddedMCPContext) + if !ok { + return nil + } + server := mcp.NewServer(&mcp.Implementation{ + Name: h.config.AgentName, Version: h.config.AgentVersion, + }, nil) + for _, descriptor := range requestContext.surface.All() { + descriptor := descriptor + server.AddTool(&mcp.Tool{ + Name: descriptor.Name, Description: descriptor.Description, + InputSchema: descriptor.InputSchema, OutputSchema: descriptor.OutputSchema, + }, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + result, err := RunCallback(ctx, h.config.Runner, func(callCtx context.Context) (json.RawMessage, error) { + return h.config.Invoke(callCtx, requestContext.session, descriptor.Name, req.Params.Arguments) + }) + if err != nil { + requestContext.failure.mu.Lock() + requestContext.failure.message = callbackMessage(err) + requestContext.failure.mu.Unlock() + return mcpError(callbackMessage(err)), nil + } + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: string(result)}}, + StructuredContent: result, + }, nil + }) + } + return server +} + +func requestID(body []byte) json.RawMessage { + var request struct { + ID json.RawMessage `json:"id"` + } + if json.Unmarshal(body, &request) != nil || len(request.ID) == 0 || !json.Valid(request.ID) { + return json.RawMessage("null") + } + return append(json.RawMessage(nil), request.ID...) +} + +func writeMCPCallbackError(w http.ResponseWriter, id json.RawMessage, err error) { + writeMCPEnvelope(w, id, callbackMessage(err)) +} + +func callbackMessage(err error) string { + switch { + case errors.Is(err, ErrCallbackCapacity): + return "host callback capacity exceeded" + case errors.Is(err, context.DeadlineExceeded): + return "host callback timed out" + case errors.Is(err, context.Canceled): + return "host callback canceled" + default: + return "host callback failed" + } +} + +func writeMCPEnvelope(w http.ResponseWriter, id json.RawMessage, message string) { + if len(id) == 0 || !json.Valid(id) { + id = json.RawMessage("null") + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Error struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + }{JSONRPC: "2.0", ID: id, Error: struct { + Code int `json:"code"` + Message string `json:"message"` + }{Code: -32603, Message: message}}) +} + +func authorizationStatus(err error) int { + var statusError interface{ HTTPStatus() int } + if errors.As(err, &statusError) { + status := statusError.HTTPStatus() + if status == http.StatusUnauthorized || status == http.StatusForbidden { + return status + } + } + return 0 +} + +// mcpError creates an MCP tool error result shared by standalone and embedded servers. +func mcpError(msg string) *mcp.CallToolResult { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: msg}}, + IsError: true, + } +} diff --git a/internal/apiserver/embedded_mcp_test.go b/internal/apiserver/embedded_mcp_test.go new file mode 100644 index 000000000..d2ac50973 --- /dev/null +++ b/internal/apiserver/embedded_mcp_test.go @@ -0,0 +1,414 @@ +package apiserver + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sort" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +type embeddedTestSession struct{ name string } + +type embeddedTestHost struct { + resolve func(context.Context, *http.Request) (any, error) + tools func(context.Context, any) ([]ToolDescriptor, error) + invoke func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) +} + +func embeddedHandler(t *testing.T, host embeddedTestHost, timeout time.Duration, capacity int) http.Handler { + t.Helper() + runner, err := NewCallbackRunner(timeout, capacity) + if err != nil { + t.Fatal(err) + } + return NewEmbeddedMCPHandler(EmbeddedMCPConfig{ + AgentName: "embedded-test", AgentVersion: "1", Runner: runner, + Resolve: host.resolve, Tools: host.tools, Invoke: host.invoke, + }) +} + +func objectTool(name string) ToolDescriptor { + return ToolDescriptor{Name: name, Description: name, InputSchema: json.RawMessage(`{"type":"object"}`)} +} + +func mcpPost(handler http.Handler, body string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/mcp/", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "application/json, text/event-stream") + handler.ServeHTTP(recorder, request) + return recorder +} + +func mcpResult(t *testing.T, recorder *httptest.ResponseRecorder) map[string]any { + t.Helper() + body := recorder.Body.String() + for _, line := range strings.Split(body, "\n") { + if strings.HasPrefix(line, "data: ") { + body = strings.TrimPrefix(line, "data: ") + break + } + } + var result map[string]any + if err := json.Unmarshal([]byte(body), &result); err != nil { + t.Fatalf("response is not JSON: %v; body=%q", err, recorder.Body.String()) + } + return result +} + +func embeddedToolNames(t *testing.T, result map[string]any) []string { + t.Helper() + value, ok := result["result"].(map[string]any) + if !ok { + t.Fatalf("missing result: %#v", result) + } + items, ok := value["tools"].([]any) + if !ok { + t.Fatalf("missing tools: %#v", result) + } + names := make([]string, 0, len(items)) + for _, item := range items { + names = append(names, item.(map[string]any)["name"].(string)) + } + sort.Strings(names) + return names +} + +func listRequest(id int) string { + return `{"jsonrpc":"2.0","id":` + jsonNumber(id) + `,"method":"tools/list","params":{}}` +} + +func jsonNumber(n int) string { + b, _ := json.Marshal(n) + return string(b) +} + +func TestEmbeddedMCPExactRequestLocalSurfaces(t *testing.T) { + sessionA, sessionB := &embeddedTestSession{"A"}, &embeddedTestSession{"B"} + host := embeddedTestHost{ + resolve: func(_ context.Context, r *http.Request) (any, error) { + if r.Header.Get("X-Session") == "B" { + return sessionB, nil + } + return sessionA, nil + }, + tools: func(_ context.Context, session any) ([]ToolDescriptor, error) { + if session == sessionB { + return []ToolDescriptor{objectTool("shared"), objectTool("beta_only")}, nil + } + return []ToolDescriptor{objectTool("shared"), objectTool("alpha_only")}, nil + }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { return nil, nil }, + } + handler := embeddedHandler(t, host, time.Second, 8) + request := func(session string) []string { + r := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/mcp/", strings.NewReader(listRequest(1))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("X-Session", session) + handler.ServeHTTP(r, req) + return embeddedToolNames(t, mcpResult(t, r)) + } + a, b := request("A"), request("B") + if len(a) == 0 || len(b) == 0 { + t.Fatal("instrument failure") + } + if strings.Join(a, ",") != "alpha_only,shared" || strings.Join(b, ",") != "beta_only,shared" { + t.Fatalf("surfaces A=%v B=%v", a, b) + } + if strings.Contains(strings.Join(a, ","), "beta_only") || strings.Contains(strings.Join(b, ","), "alpha_only") { + t.Fatal("foreign sentinel leaked") + } +} + +func TestEmbeddedMCPDispatchAuthorizationAndSessionIdentity(t *testing.T) { + sessionA, sessionB := &embeddedTestSession{"A"}, &embeddedTestSession{"B"} + var calls atomic.Int32 + var gotSession any + var gotArgs json.RawMessage + host := embeddedTestHost{ + resolve: func(_ context.Context, r *http.Request) (any, error) { + if r.Header.Get("X-Session") == "B" { + return sessionB, nil + } + return sessionA, nil + }, + tools: func(_ context.Context, session any) ([]ToolDescriptor, error) { + if session == sessionA { + return []ToolDescriptor{objectTool("alpha_only")}, nil + } + return []ToolDescriptor{objectTool("beta_only")}, nil + }, + invoke: func(_ context.Context, session any, _ string, args json.RawMessage) (json.RawMessage, error) { + calls.Add(1) + gotSession = session + gotArgs = append(json.RawMessage(nil), args...) + return json.RawMessage(`{"ok":true}`), nil + }, + } + handler := embeddedHandler(t, host, time.Second, 8) + call := `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"alpha_only","arguments":{"nonce":"A-137"}}}` + request := func(session string) map[string]any { + r := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/mcp/", strings.NewReader(call)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("X-Session", session) + handler.ServeHTTP(r, req) + return mcpResult(t, r) + } + request("A") + if calls.Load() != 1 || gotSession != sessionA || gotSession == sessionB || !bytes.Equal(gotArgs, []byte(`{"nonce":"A-137"}`)) { + t.Fatalf("dispatch calls=%d session=%p args=%s", calls.Load(), gotSession, gotArgs) + } + calls.Store(0) + result := request("B") + if calls.Load() != 0 || result["error"] == nil { + t.Fatalf("unauthorized dispatch result=%#v calls=%d", result, calls.Load()) + } +} + +func TestEmbeddedMCPSubmitFeedbackIsCallerSupplied(t *testing.T) { + var enabled atomic.Bool + var calls atomic.Int32 + host := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { return "session", nil }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { + if enabled.Load() { + return []ToolDescriptor{objectTool("submit_feedback")}, nil + } + return nil, nil + }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { + calls.Add(1) + return json.RawMessage(`{"ok":true}`), nil + }, + } + handler := embeddedHandler(t, host, time.Second, 8) + if names := embeddedToolNames(t, mcpResult(t, mcpPost(handler, listRequest(1)))); len(names) != 0 { + t.Fatalf("ambient tools: %v", names) + } + enabled.Store(true) + if names := embeddedToolNames(t, mcpResult(t, mcpPost(handler, listRequest(2)))); strings.Join(names, ",") != "submit_feedback" { + t.Fatalf("supplied tools: %v", names) + } + mcpPost(handler, `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"submit_feedback","arguments":{}}}`) + if calls.Load() != 1 { + t.Fatalf("invoke calls=%d", calls.Load()) + } +} + +func TestEmbeddedMCPStatelessTransport(t *testing.T) { + host := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { return "s", nil }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { + return []ToolDescriptor{objectTool("sentinel")}, nil + }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { return nil, nil }, + } + handler := embeddedHandler(t, host, time.Second, 8) + get := httptest.NewRecorder() + handler.ServeHTTP(get, httptest.NewRequest(http.MethodGet, "/mcp/", nil)) + if get.Code != http.StatusMethodNotAllowed || get.Header().Get("Allow") != "POST" || strings.TrimSpace(get.Body.String()) != "Method Not Allowed" { + t.Fatalf("GET status=%d allow=%q body=%q", get.Code, get.Header().Get("Allow"), get.Body.String()) + } + post := mcpPost(handler, listRequest(4)) + if post.Code != http.StatusOK || !strings.HasPrefix(post.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(post.Body.String(), "event: message") { + t.Fatalf("POST status=%d content-type=%q body=%q", post.Code, post.Header().Get("Content-Type"), post.Body.String()) + } + if names := embeddedToolNames(t, mcpResult(t, post)); strings.Join(names, ",") != "sentinel" { + t.Fatalf("tools=%v", names) + } +} + +func TestEmbeddedMCPPanicSafetyAndDescriptorValidation(t *testing.T) { + var schema atomic.Value + schema.Store(json.RawMessage(`{"type":"object"}`)) + var calls atomic.Int32 + host := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { return "s", nil }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { + return []ToolDescriptor{{Name: "tool", InputSchema: schema.Load().(json.RawMessage)}}, nil + }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { + calls.Add(1) + return json.RawMessage(`{"ok":true}`), nil + }, + } + handler := embeddedHandler(t, host, time.Second, 8) + invalid := []struct { + schema json.RawMessage + wantMsg string + }{ + {nil, "input schema is required"}, + {json.RawMessage(`"scalar"`), "invalid input schema"}, + {json.RawMessage(`{"type":"object","properties":{"bad":{"type":"array","x-mcp-header":"Bad Header"}}}`), "invalid parameter header annotations"}, + } + for i, tc := range invalid { + schema.Store(tc.schema) + r := mcpPost(handler, listRequest(10+i)) + result := mcpResult(t, r) + if r.Code != http.StatusOK || result["error"] == nil || calls.Load() != 0 { + t.Fatalf("invalid %d: status=%d result=%#v calls=%d", i, r.Code, result, calls.Load()) + } + // The message must be the gateway's LOUD rejection, not the recover() + // backstop ("host tool registration failed") — the backstop alone would + // silently accept a descriptor class the contract calls invalid. + message, _ := result["error"].(map[string]any)["message"].(string) + if !strings.Contains(message, tc.wantMsg) || strings.Contains(message, "host tool registration failed") { + t.Fatalf("invalid %d: message=%q want substring %q from gateway validation", i, message, tc.wantMsg) + } + } + schema.Store(json.RawMessage(`{"type":"object"}`)) + good := mcpPost(handler, `{"jsonrpc":"2.0","id":20,"method":"tools/call","params":{"name":"tool","arguments":{}}}`) + if good.Code != http.StatusOK || mcpResult(t, good)["result"] == nil || calls.Load() != 1 { + t.Fatalf("good status=%d body=%q calls=%d", good.Code, good.Body.String(), calls.Load()) + } +} + +type testAuthorizationError struct{ status int } + +func (e testAuthorizationError) Error() string { return "denied" } +func (e testAuthorizationError) HTTPStatus() int { return e.status } + +func TestEmbeddedMCPFrozenCallbackEnvelopes(t *testing.T) { + block := func(ctx context.Context) error { <-ctx.Done(); return ctx.Err() } + for _, stage := range []string{"resolve", "tools", "invoke"} { + bodies := []string{ + `{"jsonrpc":"2.0","id":"echo-me","method":"tools/call","params":{"name":"tool","arguments":{}}}`, + } + if stage == "resolve" { + bodies = append(bodies, `not-json`) + } + for _, body := range bodies { + var next atomic.Int32 + host := embeddedTestHost{ + resolve: func(ctx context.Context, _ *http.Request) (any, error) { + if stage == "resolve" { + return nil, block(ctx) + } + next.Add(1) + return "s", nil + }, + tools: func(ctx context.Context, _ any) ([]ToolDescriptor, error) { + if stage == "tools" { + return nil, block(ctx) + } + next.Add(1) + return []ToolDescriptor{objectTool("tool")}, nil + }, + invoke: func(ctx context.Context, _ any, _ string, _ json.RawMessage) (json.RawMessage, error) { + if stage == "invoke" { + return nil, block(ctx) + } + next.Add(1) + return nil, nil + }, + } + start := time.Now() + r := mcpPost(embeddedHandler(t, host, 20*time.Millisecond, 4), body) + elapsed := time.Since(start) + result := mcpResult(t, r) + protocolError := result["error"].(map[string]any) + if r.Code == 500 || r.Code == 504 || r.Code != 200 || r.Header().Get("Content-Type") != "application/json" || protocolError["code"] != float64(-32603) || protocolError["message"] != "host callback timed out" || elapsed < 20*time.Millisecond || elapsed >= 250*time.Millisecond { + t.Fatalf("stage=%s status=%d elapsed=%v result=%#v", stage, r.Code, elapsed, result) + } + if body == "not-json" && result["id"] != nil { + t.Fatalf("malformed id=%#v", result["id"]) + } + if body != "not-json" && result["id"] != "echo-me" { + t.Fatalf("echoed id=%#v", result["id"]) + } + wantNext := int32(0) + if stage == "tools" { + wantNext = 1 + } + if stage == "invoke" { + wantNext = 2 + } + if next.Load() != wantNext { + t.Fatalf("stage=%s next=%d want=%d", stage, next.Load(), wantNext) + } + } + } + canceled := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { return nil, context.Canceled }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { return nil, nil }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { return nil, nil }, + } + if message := mcpResult(t, mcpPost(embeddedHandler(t, canceled, time.Second, 4), listRequest(1)))["error"].(map[string]any)["message"]; message != "host callback canceled" { + t.Fatalf("message=%v", message) + } + for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden} { + auth := canceled + auth.resolve = func(context.Context, *http.Request) (any, error) { return nil, testAuthorizationError{status} } + r := mcpPost(embeddedHandler(t, auth, time.Second, 4), listRequest(1)) + if r.Code != status || r.Code == 500 || r.Code == 504 { + t.Fatalf("authorization status=%d want=%d", r.Code, status) + } + } +} + +func TestEmbeddedMCPOverloadEnvelopeAndFastControl(t *testing.T) { + release := make(chan struct{}) + blocking := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { <-release; return "s", nil }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { return nil, nil }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { return nil, nil }, + } + handler := embeddedHandler(t, blocking, 30*time.Millisecond, 4) + var overload, badStatus atomic.Int32 + var wg sync.WaitGroup + for i := 0; i < 40; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + recorder := mcpPost(handler, listRequest(id)) + if recorder.Code != http.StatusOK { + badStatus.Add(1) + } + result := mcpResult(t, recorder) + if e, ok := result["error"].(map[string]any); ok && e["message"] == "host callback capacity exceeded" && e["code"] == float64(-32603) { + overload.Add(1) + } + }(i) + } + wg.Wait() + if overload.Load() == 0 || badStatus.Load() != 0 { + t.Fatalf("overload=%d badStatus=%d", overload.Load(), badStatus.Load()) + } + close(release) + fast := embeddedTestHost{ + resolve: func(context.Context, *http.Request) (any, error) { return "s", nil }, + tools: func(context.Context, any) ([]ToolDescriptor, error) { return nil, nil }, + invoke: func(context.Context, any, string, json.RawMessage) (json.RawMessage, error) { + return nil, errors.New("unused") + }, + } + fastHandler := embeddedHandler(t, fast, time.Second, 4) + overload.Store(0) + wg = sync.WaitGroup{} + for i := 0; i < 40; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + result := mcpResult(t, mcpPost(fastHandler, listRequest(id))) + if e, ok := result["error"].(map[string]any); ok && e["message"] == "host callback capacity exceeded" { + overload.Add(1) + } + }(i) + } + wg.Wait() + if overload.Load() != 0 { + t.Fatalf("fast overloads=%d", overload.Load()) + } +} diff --git a/internal/apiserver/mcp.go b/internal/apiserver/mcp.go index 2ba08c176..d12a549a9 100644 --- a/internal/apiserver/mcp.go +++ b/internal/apiserver/mcp.go @@ -374,13 +374,3 @@ func ailangTypeToJSONSchema(ailangType string) string { return "string" } } - -// mcpError creates an MCP error result. -func mcpError(msg string) *mcp.CallToolResult { - return &mcp.CallToolResult{ - Content: []mcp.Content{ - &mcp.TextContent{Text: msg}, - }, - IsError: true, - } -} diff --git a/serveapi/serveapi.go b/serveapi/serveapi.go index 776db7d5f..438eb8811 100644 --- a/serveapi/serveapi.go +++ b/serveapi/serveapi.go @@ -58,6 +58,22 @@ type AgentInfo struct { Version string } +// AuthorizationError maps resolver rejection to an HTTP 401 or 403 response. +type AuthorizationError struct { + Status int + Err error +} + +func (e *AuthorizationError) Error() string { + if e.Err != nil { + return e.Err.Error() + } + return http.StatusText(e.Status) +} + +func (e *AuthorizationError) Unwrap() error { return e.Err } +func (e *AuthorizationError) HTTPStatus() int { return e.Status } + type Config struct { Resolver SessionResolver Tools ToolSource @@ -70,8 +86,9 @@ type Config struct { // Server is an embeddable protocol surface. Wire adapters are completed in // later milestones; M1 establishes and validates the host contract. type Server struct { - config Config - runner *apiserver.CallbackRunner + config Config + runner *apiserver.CallbackRunner + mcpHandler http.Handler } func New(cfg Config) (*Server, error) { @@ -106,11 +123,33 @@ func New(cfg Config) (*Server, error) { if err != nil { return nil, fmt.Errorf("configure callback runner: %w", err) } - return &Server{config: cfg, runner: runner}, nil + s := &Server{config: cfg, runner: runner} + s.mcpHandler = apiserver.NewEmbeddedMCPHandler(apiserver.EmbeddedMCPConfig{ + AgentName: cfg.Agent.Name, AgentVersion: cfg.Agent.Version, Runner: runner, + Resolve: func(ctx context.Context, request *http.Request) (any, error) { + return cfg.Resolver.ResolveSession(ctx, request) + }, + Tools: func(ctx context.Context, session any) ([]apiserver.ToolDescriptor, error) { + descriptors, err := cfg.Tools.Tools(ctx, session) + if err != nil { + return nil, err + } + result := make([]apiserver.ToolDescriptor, len(descriptors)) + for i, descriptor := range descriptors { + result[i] = apiserver.ToolDescriptor(descriptor) + } + return result, nil + }, + Invoke: func(ctx context.Context, session any, name string, arguments json.RawMessage) (json.RawMessage, error) { + result, err := cfg.Invoker.Invoke(ctx, session, Invocation{Name: name, Arguments: arguments}) + return result.Value, err + }, + }) + return s, nil } -// MCPHandler returns the MCP endpoint handler. Wire behavior lands in M2. -func (s *Server) MCPHandler() http.Handler { return http.NotFoundHandler() } +// MCPHandler returns the request-scoped stateless MCP endpoint handler. +func (s *Server) MCPHandler() http.Handler { return s.mcpHandler } // A2AHandler returns the A2A endpoint handler. Wire behavior lands in M3. func (s *Server) A2AHandler() http.Handler { return http.NotFoundHandler() }