fix(contrib/mark3labs/mcp-go): preserve raw input schema fields on intent injection (#4946)

## Summary

`mcp.ToolInputSchema` only models `type`, `properties`, `required`, and `$defs`. Tools registered via `NewToolWithRawSchema` may declare other JSON Schema keywords (`additionalProperties`, `oneOf`, `patternProperties`, `pattern`, `enum`, `minimum`, `$schema`, …) — the previous round-trip through `ToolInputSchema` dropped all of them, so enabling intent capture quietly relaxed every raw-schema tool.

Fix: inject telemetry into the raw JSON as a generic `map[string]any` and re-marshal. Unknown keys pass through verbatim.

Same bug exists in the dd-source copy (not fixed there). Found by codex-connector on #4939.

## Test plan

- [x] `TestIntentCaptureRawInputSchemaPreservesUnknownFields` — raw schema with `additionalProperties: false` and `oneOf` survives injection.
- [x] All existing intent-capture tests still pass.

Stacked on #4940.

Co-authored-by: jboolean <julian.boilen@datadoghq.com>

Author: gh-worker-dd-mergequeue-cf854d[bot]

contrib/mark3labs/mcp-go/intent_capture.go

@@ -44,13 +44,19 @@ func injectTelemetryListToolsHook(ctx context.Context, id any, message *mcp.List
 	for i := range result.Tools {
 		t := &result.Tools[i]
 
+		// mcp.ToolInputSchema only models type/properties/required/$defs;
+		// for tools defined with NewToolWithRawSchema we mutate the raw
+		// JSON via a generic map so keywords like additionalProperties,
+		// oneOf, or patternProperties pass through verbatim.
 		if t.RawInputSchema != nil {
-			var schema mcp.ToolInputSchema
-			if err := json.Unmarshal(t.RawInputSchema, &schema); err != nil {
-				continue
+			if newRaw, ok := injectTelemetryIntoRawSchema(t.RawInputSchema); ok {
+				t.RawInputSchema = newRaw
+				// mcp.NewTool sets InputSchema.Type="object" by default; keeping it
+				// alongside RawInputSchema makes Tool.MarshalJSON return a schema-
+				// conflict error.
+				t.InputSchema = mcp.ToolInputSchema{}
 			}
-			t.InputSchema = schema
-			t.RawInputSchema = nil
+			continue
 		}
 
 		if t.InputSchema.Type == "" {
@@ -72,6 +78,38 @@ func injectTelemetryListToolsHook(ctx context.Context, id any, message *mcp.List
 	}
 }
 
+// injectTelemetryIntoRawSchema mutates a raw JSON Schema document to add the
+// telemetry property and require it. Unknown top-level keywords pass through
+// verbatim. Returns false when the input can't be parsed as a JSON object.
+func injectTelemetryIntoRawSchema(raw json.RawMessage) (json.RawMessage, bool) {
+	var schema map[string]any
+	if json.Unmarshal(raw, &schema) != nil || schema == nil {
+		return nil, false
+	}
+	if _, ok := schema["type"]; !ok {
+		schema["type"] = "object"
+	}
+	props, _ := schema["properties"].(map[string]any)
+	props = maps.Clone(props)
+	if props == nil {
+		props = map[string]any{}
+	}
+	props[instrmcp.TelemetryKey] = telemetrySchema()
+	schema["properties"] = props
+
+	required, _ := schema["required"].([]any)
+	if !slices.Contains(required, any(instrmcp.TelemetryKey)) {
+		required = append(slices.Clone(required), instrmcp.TelemetryKey)
+	}
+	schema["required"] = required
+
+	out, err := json.Marshal(schema)
+	if err != nil {
+		return nil, false
+	}
+	return out, true
+}
+
 // Removing tracing parameters from the tool call request so its not sent to the tool.
 // This must be registered after the tool handler middleware (mcp-go runs middleware in registration order).
 // This removes the telemetry parameter before user-defined middleware or tool handlers can see it.

contrib/mark3labs/mcp-go/intent_capture_test.go

@@ -92,6 +92,61 @@ func TestIntentCapture(t *testing.T) {
 	assert.Equal(t, "test intent description", toolSpan.Meta["intent"])
 }
 
+func TestIntentCaptureRawInputSchemaViaNewToolListsWithoutConflict(t *testing.T) {
+	// mcp.NewTool defaults InputSchema.Type to "object"; combined with
+	// WithRawInputSchema this leaves BOTH set, and Tool.MarshalJSON refuses
+	// to encode a tool with both. Intent capture must clear the structured
+	// schema when it keeps the raw one.
+	tt := testTracer(t)
+	defer tt.Stop()
+
+	srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{IntentCaptureEnabled: true}))
+	srv.AddTool(mcp.NewTool("raw_tool",
+		mcp.WithDescription("raw"),
+		mcp.WithRawInputSchema(json.RawMessage(`{"type":"object","properties":{"q":{"type":"string"}},"required":["q"]}`)),
+	), func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultText("ok"), nil
+	})
+
+	listResp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`))
+	var listResult map[string]interface{}
+	require.NoError(t, json.Unmarshal(mustMarshal(listResp), &listResult))
+
+	require.NotNil(t, listResult["result"], "tools/list returned error: %v", listResult["error"])
+	tools := listResult["result"].(map[string]interface{})["tools"].([]interface{})
+	require.Len(t, tools, 1)
+	props := tools[0].(map[string]interface{})["inputSchema"].(map[string]interface{})["properties"].(map[string]interface{})
+	assert.Contains(t, props, "telemetry")
+}
+
+func TestIntentCaptureRawInputSchemaPreservesUnknownFields(t *testing.T) {
+	// mcp.ToolInputSchema doesn't model additionalProperties/oneOf/etc;
+	// intent capture must not silently strip those when injecting telemetry.
+	tt := testTracer(t)
+	defer tt.Stop()
+
+	srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{IntentCaptureEnabled: true}))
+	srv.AddTool(mcp.NewToolWithRawSchema("raw_tool", "raw", json.RawMessage(`{
+		"type": "object",
+		"properties": {"app_id": {"type": "string"}},
+		"required": ["app_id"],
+		"additionalProperties": false,
+		"oneOf": [{"required": ["app_id"]}]
+	}`)), func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultText("ok"), nil
+	})
+
+	listResp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`))
+	var listResult map[string]interface{}
+	require.NoError(t, json.Unmarshal(mustMarshal(listResp), &listResult))
+	schema := listResult["result"].(map[string]interface{})["tools"].([]interface{})[0].(map[string]interface{})["inputSchema"].(map[string]interface{})
+
+	assert.Contains(t, schema["properties"].(map[string]interface{}), "telemetry")
+	assert.Contains(t, schema["required"].([]interface{}), "telemetry")
+	assert.Equal(t, false, schema["additionalProperties"])
+	assert.NotNil(t, schema["oneOf"])
+}
+
 func TestIntentCaptureRawInputSchema(t *testing.T) {
 	tt := testTracer(t)
 	defer tt.Stop()