Skip to content

Commit 831b389

Browse files
ChrisJBurnsChrisJBurns
authored
Add Custom OpenTelemetry Resource Attributes Support (#2285)
* otel custom attributes Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com> * log Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com> * lint Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com> * removes custom attribute docs Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com> * updates docs Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com> --------- Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com>
1 parent 406d888 commit 831b389

File tree

10 files changed

+357
-9
lines changed

10 files changed

+357
-9
lines changed

‎cmd/thv/app/run_flags.go‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/stacklok/toolhive/pkg/container/runtime"
1616
"github.com/stacklok/toolhive/pkg/environment"
1717
"github.com/stacklok/toolhive/pkg/ignore"
18+
"github.com/stacklok/toolhive/pkg/logger"
1819
"github.com/stacklok/toolhive/pkg/networking"
1920
"github.com/stacklok/toolhive/pkg/oauth"
2021
"github.com/stacklok/toolhive/pkg/process"
@@ -79,6 +80,7 @@ type RunFlags struct {
7980
OtelInsecure bool
8081
OtelEnablePrometheusMetricsPath bool
8182
OtelEnvironmentVariables []string // renamed binding to otel-env-vars
83+
OtelCustomAttributes string // Custom attributes in key=value format
8284

8385
// Network isolation
8486
IsolateNetwork bool
@@ -199,6 +201,8 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) {
199201
"Enable Prometheus-style /metrics endpoint on the main transport port")
200202
cmd.Flags().StringArrayVar(&config.OtelEnvironmentVariables, "otel-env-vars", nil,
201203
"Environment variable names to include in OpenTelemetry spans (comma-separated: ENV1,ENV2)")
204+
cmd.Flags().StringVar(&config.OtelCustomAttributes, "otel-custom-attributes", "",
205+
"Custom resource attributes for OpenTelemetry in key=value format (e.g., server_type=prod,region=us-east-1,team=platform)")
202206

203207
cmd.Flags().BoolVar(&config.IsolateNetwork, "isolate-network", false,
204208
"Isolate the container network from the host (default: false)")
@@ -321,7 +325,7 @@ func setupTelemetryConfiguration(cmd *cobra.Command, runFlags *RunFlags) *teleme
321325

322326
return createTelemetryConfig(finalOtelEndpoint, finalOtelEnablePrometheusMetricsPath,
323327
runFlags.OtelServiceName, runFlags.OtelTracingEnabled, runFlags.OtelMetricsEnabled, finalOtelSamplingRate,
324-
runFlags.OtelHeaders, finalOtelInsecure, finalOtelEnvironmentVariables)
328+
runFlags.OtelHeaders, finalOtelInsecure, finalOtelEnvironmentVariables, runFlags.OtelCustomAttributes)
325329
}
326330

327331
// setupRuntimeAndValidation creates container runtime and selects environment variable validator
@@ -779,7 +783,7 @@ func createOIDCConfig(oidcIssuer, oidcAudience, oidcJwksURL, oidcIntrospectionUR
779783
// createTelemetryConfig creates a telemetry configuration if any telemetry parameters are provided
780784
func createTelemetryConfig(otelEndpoint string, otelEnablePrometheusMetricsPath bool,
781785
otelServiceName string, otelTracingEnabled bool, otelMetricsEnabled bool, otelSamplingRate float64, otelHeaders []string,
782-
otelInsecure bool, otelEnvironmentVariables []string) *telemetry.Config {
786+
otelInsecure bool, otelEnvironmentVariables []string, otelCustomAttributes string) *telemetry.Config {
783787
if otelEndpoint == "" && !otelEnablePrometheusMetricsPath {
784788
return nil
785789
}
@@ -812,6 +816,14 @@ func createTelemetryConfig(otelEndpoint string, otelEnablePrometheusMetricsPath
812816
}
813817
}
814818

819+
// Parse custom attributes
820+
customAttrs, err := telemetry.ParseCustomAttributes(otelCustomAttributes)
821+
if err != nil {
822+
// Log the error but don't fail - telemetry is optional
823+
logger.Warnf("Failed to parse custom attributes: %v", err)
824+
customAttrs = nil
825+
}
826+
815827
return &telemetry.Config{
816828
Endpoint: otelEndpoint,
817829
ServiceName: serviceName,
@@ -823,6 +835,7 @@ func createTelemetryConfig(otelEndpoint string, otelEnablePrometheusMetricsPath
823835
Insecure: otelInsecure,
824836
EnablePrometheusMetricsPath: otelEnablePrometheusMetricsPath,
825837
EnvironmentVariables: processedEnvVars,
838+
CustomAttributes: customAttrs,
826839
}
827840
}
828841

‎docs/cli/thv_run.md‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/server/docs.go‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/server/swagger.json‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/server/swagger.yaml‎

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/runner/config_builder.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ func WithOIDCConfig(
353353
}
354354
}
355355

356-
// WithTelemetryConfig configures telemetry settings
356+
// WithTelemetryConfig configures telemetry settings (legacy - custom attributes handled via middleware)
357357
func WithTelemetryConfig(
358358
otelEndpoint string,
359359
otelEnablePrometheusMetricsPath bool,

‎pkg/telemetry/attributes.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Package telemetry provides OpenTelemetry instrumentation for ToolHive MCP server proxies.
2+
package telemetry
3+
4+
import (
5+
"fmt"
6+
"strings"
7+
8+
"go.opentelemetry.io/otel/attribute"
9+
)
10+
11+
// ParseCustomAttributes parses a comma-separated list of key=value pairs into a map.
12+
// Example input: "server_type=production,region=us-east-1,team=platform"
13+
// Returns a map[string]string that can be converted to resource attributes.
14+
func ParseCustomAttributes(input string) (map[string]string, error) {
15+
if input == "" {
16+
return map[string]string{}, nil
17+
}
18+
19+
attributes := make(map[string]string)
20+
pairs := strings.Split(input, ",")
21+
22+
for _, pair := range pairs {
23+
trimmedPair := strings.TrimSpace(pair)
24+
if trimmedPair == "" {
25+
continue
26+
}
27+
28+
parts := strings.SplitN(trimmedPair, "=", 2)
29+
if len(parts) != 2 {
30+
return nil, fmt.Errorf("invalid attribute format '%s': expected key=value", trimmedPair)
31+
}
32+
33+
key := strings.TrimSpace(parts[0])
34+
value := strings.TrimSpace(parts[1])
35+
36+
if key == "" {
37+
return nil, fmt.Errorf("empty attribute key in '%s'", trimmedPair)
38+
}
39+
40+
// Store as key-value pair in map
41+
attributes[key] = value
42+
}
43+
44+
return attributes, nil
45+
}
46+
47+
// ConvertMapToAttributes converts a map[string]string to OpenTelemetry attributes
48+
func ConvertMapToAttributes(attrs map[string]string) []attribute.KeyValue {
49+
var result []attribute.KeyValue
50+
for k, v := range attrs {
51+
result = append(result, attribute.String(k, v))
52+
}
53+
return result
54+
}

‎pkg/telemetry/attributes_test.go‎

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
package telemetry
2+
3+
import (
4+
"testing"
5+
6+
"go.opentelemetry.io/otel/attribute"
7+
)
8+
9+
func TestParseCustomAttributes(t *testing.T) {
10+
t.Parallel()
11+
tests := []struct {
12+
name string
13+
input string
14+
want []attribute.KeyValue
15+
wantErr bool
16+
}{
17+
{
18+
name: "single attribute",
19+
input: "environment=production",
20+
want: []attribute.KeyValue{
21+
attribute.String("environment", "production"),
22+
},
23+
},
24+
{
25+
name: "multiple attributes",
26+
input: "environment=production,region=us-east-1,team=platform",
27+
want: []attribute.KeyValue{
28+
attribute.String("environment", "production"),
29+
attribute.String("region", "us-east-1"),
30+
attribute.String("team", "platform"),
31+
},
32+
},
33+
{
34+
name: "attributes with spaces",
35+
input: " environment = production , region = us-east-1 ",
36+
want: []attribute.KeyValue{
37+
attribute.String("environment", "production"),
38+
attribute.String("region", "us-east-1"),
39+
},
40+
},
41+
{
42+
name: "attribute with special characters",
43+
input: "service.name=my-service,service.version=1.2.3",
44+
want: []attribute.KeyValue{
45+
attribute.String("service.name", "my-service"),
46+
attribute.String("service.version", "1.2.3"),
47+
},
48+
},
49+
{
50+
name: "attribute with underscore",
51+
input: "server_type=production,deployment_id=12345",
52+
want: []attribute.KeyValue{
53+
attribute.String("server_type", "production"),
54+
attribute.String("deployment_id", "12345"),
55+
},
56+
},
57+
{
58+
name: "empty input",
59+
input: "",
60+
want: []attribute.KeyValue{},
61+
},
62+
{
63+
name: "trailing comma",
64+
input: "environment=production,",
65+
want: []attribute.KeyValue{
66+
attribute.String("environment", "production"),
67+
},
68+
},
69+
{
70+
name: "multiple equals in value",
71+
input: "url=https://example.com/path?query=value",
72+
want: []attribute.KeyValue{
73+
attribute.String("url", "https://example.com/path?query=value"),
74+
},
75+
},
76+
{
77+
name: "missing equals",
78+
input: "environment",
79+
wantErr: true,
80+
},
81+
{
82+
name: "missing key",
83+
input: "=production",
84+
wantErr: true,
85+
},
86+
{
87+
name: "empty key",
88+
input: " =production",
89+
wantErr: true,
90+
},
91+
{
92+
name: "empty value is allowed",
93+
input: "debug=",
94+
want: []attribute.KeyValue{
95+
attribute.String("debug", ""),
96+
},
97+
},
98+
{
99+
name: "mixed valid and invalid",
100+
input: "environment=production,invalid,region=us-east-1",
101+
wantErr: true,
102+
},
103+
{
104+
name: "numeric-like values as strings",
105+
input: "port=8080,count=100,ratio=0.95",
106+
want: []attribute.KeyValue{
107+
attribute.String("port", "8080"),
108+
attribute.String("count", "100"),
109+
attribute.String("ratio", "0.95"),
110+
},
111+
},
112+
{
113+
name: "boolean-like values as strings",
114+
input: "enabled=true,debug=false",
115+
want: []attribute.KeyValue{
116+
attribute.String("enabled", "true"),
117+
attribute.String("debug", "false"),
118+
},
119+
},
120+
{
121+
name: "attribute with encoded characters",
122+
input: "message=Hello%20World,path=/api/v1/users",
123+
want: []attribute.KeyValue{
124+
attribute.String("message", "Hello%20World"),
125+
attribute.String("path", "/api/v1/users"),
126+
},
127+
},
128+
{
129+
name: "complex real-world example",
130+
input: "service.name=toolhive,service.namespace=default,service.instance.id=i-1234567890abcdef,cloud.provider=aws,cloud.region=us-west-2",
131+
want: []attribute.KeyValue{
132+
attribute.String("service.name", "toolhive"),
133+
attribute.String("service.namespace", "default"),
134+
attribute.String("service.instance.id", "i-1234567890abcdef"),
135+
attribute.String("cloud.provider", "aws"),
136+
attribute.String("cloud.region", "us-west-2"),
137+
},
138+
},
139+
}
140+
141+
for _, tt := range tests {
142+
tt := tt // capture range variable
143+
t.Run(tt.name, func(t *testing.T) {
144+
t.Parallel()
145+
got, err := ParseCustomAttributes(tt.input)
146+
if (err != nil) != tt.wantErr {
147+
t.Errorf("ParseCustomAttributes() error = %v, wantErr %v", err, tt.wantErr)
148+
return
149+
}
150+
if !tt.wantErr {
151+
if len(got) != len(tt.want) {
152+
t.Errorf("ParseCustomAttributes() got %d attributes, want %d", len(got), len(tt.want))
153+
return
154+
}
155+
156+
// Check that the returned map is as expected for the input.
157+
gotMap := got
158+
wantMap := make(map[string]string)
159+
for _, attr := range tt.want {
160+
wantMap[string(attr.Key)] = attr.Value.AsString()
161+
}
162+
if len(gotMap) != len(wantMap) {
163+
t.Errorf("ParseCustomAttributes() got %d attributes, want %d", len(gotMap), len(wantMap))
164+
} else {
165+
for k, v := range wantMap {
166+
if gotMap[k] != v {
167+
t.Errorf("ParseCustomAttributes()[%q] = %q, want %q", k, gotMap[k], v)
168+
}
169+
}
170+
}
171+
}
172+
})
173+
}
174+
}
175+
176+
func TestConvertMapToAttributes(t *testing.T) {
177+
t.Parallel()
178+
tests := []struct {
179+
name string
180+
input map[string]string
181+
want []attribute.KeyValue
182+
}{
183+
{
184+
name: "empty map",
185+
input: map[string]string{},
186+
want: []attribute.KeyValue{},
187+
},
188+
{
189+
name: "single attribute",
190+
input: map[string]string{"foo": "bar"},
191+
want: []attribute.KeyValue{
192+
attribute.String("foo", "bar"),
193+
},
194+
},
195+
{
196+
name: "multiple attributes",
197+
input: map[string]string{
198+
"env": "prod",
199+
"team": "platform",
200+
"release": "stable",
201+
},
202+
want: []attribute.KeyValue{
203+
attribute.String("env", "prod"),
204+
attribute.String("team", "platform"),
205+
attribute.String("release", "stable"),
206+
},
207+
},
208+
}
209+
210+
for _, tt := range tests {
211+
tt := tt // capture range variable
212+
t.Run(tt.name, func(t *testing.T) {
213+
t.Parallel()
214+
got := ConvertMapToAttributes(tt.input)
215+
if len(got) != len(tt.want) {
216+
t.Errorf("ConvertMapToAttributes() got %d attributes, want %d", len(got), len(tt.want))
217+
return
218+
}
219+
// Convert result to map for easy comparison (since map iteration order is not guaranteed)
220+
gotMap := make(map[attribute.Key]string)
221+
for _, kv := range got {
222+
gotMap[kv.Key] = kv.Value.AsString()
223+
}
224+
for _, wantKV := range tt.want {
225+
val, ok := gotMap[wantKV.Key]
226+
if !ok {
227+
t.Errorf("ConvertMapToAttributes() missing key %v", wantKV.Key)
228+
} else if val != wantKV.Value.AsString() {
229+
t.Errorf("ConvertMapToAttributes() key %v = %v, want %v", wantKV.Key, val, wantKV.Value.AsString())
230+
}
231+
}
232+
})
233+
}
234+
}

0 commit comments

Comments
 (0)