Skip to content

Commit 8bf8057

Browse files
rhnvrmrhnvrm
committed
feat: add session type labels to metrics tracking
Adds session type tracking to all tool metrics to differentiate between SSE, MCP, stdio, and unknown session types in hybrid mode. Changes: - Add session type context (sse/mcp/stdio/unknown) to all tool handlers - Extend metrics parsing to include session_type labels in Prometheus output - Add HTTP middleware to inject session type based on endpoint path (/sse → sse, /mcp → mcp) - Update all 22 tools to use context-aware tracking (previously 7 tools were missing) - Change fallback session type from "stdio" to "unknown" for better debugging - Add comprehensive tests for session type parsing Metrics now show: tool_calls_get_quotes_total{date="2025-08-06",session_type="mcp"} 45 Fixes missing metrics for POST tools and market tools that had custom handlers. All tools now properly tracked across different session types for better usage analytics.
1 parent 72a5e75 commit 8bf8057

File tree

8 files changed

+132
-45
lines changed

8 files changed

+132
-45
lines changed

‎app/app.go‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -268,10 +268,19 @@ func (app *App) createStreamableHTTPServer(mcpServer *server.MCPServer, kcManage
268268
)
269269
}
270270

271+
// withSessionType adds session type to context based on URL path
272+
func withSessionType(sessionType string, handler http.HandlerFunc) http.HandlerFunc {
273+
return func(w http.ResponseWriter, r *http.Request) {
274+
ctx := mcp.WithSessionType(r.Context(), sessionType)
275+
r = r.WithContext(ctx)
276+
handler(w, r)
277+
}
278+
}
279+
271280
// registerSSEEndpoints registers SSE-specific endpoints on the mux
272281
func (app *App) registerSSEEndpoints(mux *http.ServeMux, sse *server.SSEServer) {
273-
mux.HandleFunc("/sse", sse.ServeHTTP)
274-
mux.HandleFunc("/message", sse.ServeHTTP)
282+
mux.HandleFunc("/sse", withSessionType(mcp.SessionTypeSSE, sse.ServeHTTP))
283+
mux.HandleFunc("/message", withSessionType(mcp.SessionTypeSSE, sse.ServeHTTP))
275284
}
276285

277286
// configureAndStartServer sets up server handler and starts it
@@ -293,7 +302,7 @@ func (app *App) startHybridServer(srv *http.Server, kcManager *kc.Manager, mcpSe
293302

294303
// Register endpoints
295304
app.registerSSEEndpoints(mux, sse)
296-
mux.HandleFunc("/mcp", streamable.ServeHTTP)
305+
mux.HandleFunc("/mcp", withSessionType(mcp.SessionTypeMCP, streamable.ServeHTTP))
297306

298307
app.logger.Info("Hybrid mode enabled with both SSE and MCP endpoints on the same server")
299308
app.logger.Info("SSE endpoints available", "url", fmt.Sprintf("http://%s/sse and http://%s/message", url, url))
@@ -338,7 +347,7 @@ func (app *App) startHTTPServer(srv *http.Server, kcManager *kc.Manager, mcpServ
338347

339348
// Setup mux with common handlers
340349
mux := app.setupMux(kcManager)
341-
mux.HandleFunc("/mcp", streamable.ServeHTTP)
350+
mux.HandleFunc("/mcp", withSessionType(mcp.SessionTypeMCP, streamable.ServeHTTP))
342351

343352
app.logger.Info("MCP session manager configured with automatic cleanup for both MCP and Kite sessions")
344353
app.logger.Info("MCP Session manager configured", "session_expiry", kc.DefaultSessionDuration)

‎app/metrics/daily_metrics_test.go‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,11 +128,36 @@ func TestParseDailyMetric(t *testing.T) {
128128

129129
for _, tt := range tests {
130130
t.Run(tt.key, func(t *testing.T) {
131-
gotBase, gotDate := m.parseDailyMetric(tt.key)
131+
gotBase, _, gotDate := m.parseDailyMetric(tt.key)
132132
if gotBase != tt.expectedBase || gotDate != tt.expectedDate {
133133
t.Errorf("parseDailyMetric(%q) = (%q, %q), want (%q, %q)",
134134
tt.key, gotBase, gotDate, tt.expectedBase, tt.expectedDate)
135135
}
136136
})
137137
}
138+
139+
// Test session type parsing
140+
sessionTypeTests := []struct {
141+
key string
142+
expectedBase string
143+
expectedSession string
144+
expectedDate string
145+
}{
146+
{"tool_calls_quotes_sse_2025-08-05", "tool_calls_quotes", "sse", "2025-08-05"},
147+
{"tool_calls_quotes_mcp_2025-08-05", "tool_calls_quotes", "mcp", "2025-08-05"},
148+
{"tool_calls_quotes_stdio_2025-08-05", "tool_calls_quotes", "stdio", "2025-08-05"},
149+
{"tool_calls_quotes_unknown_2025-08-05", "tool_calls_quotes", "unknown", "2025-08-05"},
150+
{"tool_calls_quotes_2025-08-05", "tool_calls_quotes", "", "2025-08-05"}, // No session type
151+
{"tool_errors_login_session_error_sse_2025-12-31", "tool_errors_login_session_error", "sse", "2025-12-31"},
152+
}
153+
154+
for _, tt := range sessionTypeTests {
155+
t.Run(tt.key+"_session_type", func(t *testing.T) {
156+
gotBase, gotSession, gotDate := m.parseDailyMetric(tt.key)
157+
if gotBase != tt.expectedBase || gotSession != tt.expectedSession || gotDate != tt.expectedDate {
158+
t.Errorf("parseDailyMetric(%q) = (%q, %q, %q), want (%q, %q, %q)",
159+
tt.key, gotBase, gotSession, gotDate, tt.expectedBase, tt.expectedSession, tt.expectedDate)
160+
}
161+
})
162+
}
138163
}

‎app/metrics/metrics.go‎

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -214,12 +214,6 @@ func (m *Manager) isDailyMetric(key string) bool {
214214
return false
215215
}
216216

217-
// Check if we have a non-empty base name
218-
baseName := strings.Join(parts[:len(parts)-1], "_")
219-
if baseName == "" {
220-
return false
221-
}
222-
223217
// Check if the last part looks like a date (YYYY-MM-DD)
224218
lastPart := parts[len(parts)-1]
225219
if len(lastPart) != 10 || strings.Count(lastPart, "-") != 2 {
@@ -232,19 +226,38 @@ func (m *Manager) isDailyMetric(key string) bool {
232226
return false
233227
}
234228

229+
// Check if we have a non-empty base name after removing date
230+
baseName := strings.Join(parts[:len(parts)-1], "_")
231+
if baseName == "" {
232+
return false
233+
}
234+
235235
return true
236236
}
237237

238-
// parseDailyMetric extracts base name and date from a daily metric key
239-
func (m *Manager) parseDailyMetric(key string) (baseName, date string) {
238+
// parseDailyMetric extracts base name, session type and date from a daily metric key
239+
// Returns baseName, sessionType, date - sessionType will be empty if not present
240+
func (m *Manager) parseDailyMetric(key string) (baseName, sessionType, date string) {
240241
if !m.isDailyMetric(key) {
241-
return "", ""
242+
return "", "", ""
242243
}
243244

244245
parts := strings.Split(key, "_")
245246
date = parts[len(parts)-1]
247+
248+
// Check if the second-to-last part is a session type
249+
if len(parts) >= 3 {
250+
potentialSessionType := parts[len(parts)-2]
251+
if potentialSessionType == "sse" || potentialSessionType == "mcp" || potentialSessionType == "stdio" || potentialSessionType == "unknown" {
252+
sessionType = potentialSessionType
253+
baseName = strings.Join(parts[:len(parts)-2], "_")
254+
return baseName, sessionType, date
255+
}
256+
}
257+
258+
// No session type found, return base name without session type
246259
baseName = strings.Join(parts[:len(parts)-1], "_")
247-
return baseName, date
260+
return baseName, "", date
248261
}
249262

250263
// formatMetric formats a single metric in Prometheus format
@@ -278,9 +291,13 @@ func (m *Manager) WritePrometheus(buf *bytes.Buffer) {
278291

279292
// Check if this is a daily metric (has date suffix)
280293
if m.isDailyMetric(name) {
281-
baseName, date := m.parseDailyMetric(name)
294+
baseName, sessionType, date := m.parseDailyMetric(name)
282295
if baseName != "" && date != "" {
283-
m.formatMetric(buf, fmt.Sprintf("%s_total", baseName), map[string]string{"date": date}, float64(value))
296+
labels := map[string]string{"date": date}
297+
if sessionType != "" {
298+
labels["session_type"] = sessionType
299+
}
300+
m.formatMetric(buf, fmt.Sprintf("%s_total", baseName), labels, float64(value))
284301
}
285302
} else {
286303
// Regular total counter without date label

‎mcp/common.go‎

Lines changed: 47 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,34 @@ import (
1010
"github.com/zerodha/kite-mcp-server/kc"
1111
)
1212

13+
// Context key for session type
14+
type contextKey string
15+
16+
const (
17+
sessionTypeKey contextKey = "session_type"
18+
)
19+
20+
// Session type constants
21+
const (
22+
SessionTypeSSE = "sse"
23+
SessionTypeMCP = "mcp"
24+
SessionTypeStdio = "stdio"
25+
SessionTypeUnknown = "unknown"
26+
)
27+
28+
// WithSessionType adds session type to context
29+
func WithSessionType(ctx context.Context, sessionType string) context.Context {
30+
return context.WithValue(ctx, sessionTypeKey, sessionType)
31+
}
32+
33+
// SessionTypeFromContext extracts session type from context
34+
func SessionTypeFromContext(ctx context.Context) string {
35+
if sessionType, ok := ctx.Value(sessionTypeKey).(string); ok {
36+
return sessionType
37+
}
38+
return SessionTypeUnknown // default fallback for undetermined sessions
39+
}
40+
1341
// ToolHandler provides common functionality for all MCP tools
1442
type ToolHandler struct {
1543
manager *kc.Manager
@@ -20,17 +48,21 @@ func NewToolHandler(manager *kc.Manager) *ToolHandler {
2048
return &ToolHandler{manager: manager}
2149
}
2250

23-
// trackToolCall increments the daily tool usage counter
24-
func (h *ToolHandler) trackToolCall(toolName string) {
51+
// trackToolCall increments the daily tool usage counter with optional context for session type
52+
func (h *ToolHandler) trackToolCall(ctx context.Context, toolName string) {
2553
if h.manager.HasMetrics() {
26-
h.manager.IncrementDailyMetric(fmt.Sprintf("tool_calls_%s", toolName))
54+
sessionType := SessionTypeFromContext(ctx)
55+
metricName := fmt.Sprintf("tool_calls_%s_%s", toolName, sessionType)
56+
h.manager.IncrementDailyMetric(metricName)
2757
}
2858
}
2959

30-
// trackToolError increments the daily tool error counter with error type
31-
func (h *ToolHandler) trackToolError(toolName, errorType string) {
60+
// trackToolError increments the daily tool error counter with error type and optional context for session type
61+
func (h *ToolHandler) trackToolError(ctx context.Context, toolName, errorType string) {
3262
if h.manager.HasMetrics() {
33-
h.manager.IncrementDailyMetric(fmt.Sprintf("tool_errors_%s_%s", toolName, errorType))
63+
sessionType := SessionTypeFromContext(ctx)
64+
metricName := fmt.Sprintf("tool_errors_%s_%s_%s", toolName, errorType, sessionType)
65+
h.manager.IncrementDailyMetric(metricName)
3466
}
3567
}
3668

@@ -45,13 +77,13 @@ func (h *ToolHandler) WithSession(ctx context.Context, toolName string, fn func(
4577
kiteSession, isNew, err := h.manager.GetOrCreateSession(sessionID)
4678
if err != nil {
4779
h.manager.Logger.Error("Failed to establish session", "tool", toolName, "session_id", sessionID, "error", err)
48-
h.trackToolError(toolName, "session_error")
80+
h.trackToolError(ctx, toolName, "session_error")
4981
return mcp.NewToolResultError("Failed to establish a session. Please try again."), nil
5082
}
5183

5284
if isNew {
5385
h.manager.Logger.Info("New session created, login required", "tool", toolName, "session_id", sessionID)
54-
h.trackToolError(toolName, "auth_required")
86+
h.trackToolError(ctx, toolName, "auth_required")
5587
return mcp.NewToolResultError("Please log in first using the login tool"), nil
5688
}
5789

@@ -300,12 +332,12 @@ func SimpleToolHandler(manager *kc.Manager, toolName string, apiCall func(*kc.Ki
300332
handler := NewToolHandler(manager)
301333
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
302334
// Track the tool call at the handler level
303-
handler.trackToolCall(toolName)
335+
handler.trackToolCall(ctx, toolName)
304336
result, err := handler.HandleAPICall(ctx, toolName, apiCall)
305337
if err != nil {
306-
handler.trackToolError(toolName, "execution_error")
338+
handler.trackToolError(ctx, toolName, "execution_error")
307339
} else if result != nil && result.IsError {
308-
handler.trackToolError(toolName, "api_error")
340+
handler.trackToolError(ctx, toolName, "api_error")
309341
}
310342
return result, err
311343
}
@@ -316,13 +348,13 @@ func PaginatedToolHandler[T any](manager *kc.Manager, toolName string, apiCall f
316348
handler := NewToolHandler(manager)
317349
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
318350
// Track the tool call at the handler level
319-
handler.trackToolCall(toolName)
351+
handler.trackToolCall(ctx, toolName)
320352
result, err := handler.WithSession(ctx, toolName, func(session *kc.KiteSessionData) (*mcp.CallToolResult, error) {
321353
// Get the data
322354
data, err := apiCall(session)
323355
if err != nil {
324356
handler.manager.Logger.Error("API call failed", "tool", toolName, "error", err)
325-
handler.trackToolError(toolName, "api_error")
357+
handler.trackToolError(ctx, toolName, "api_error")
326358
return mcp.NewToolResultError(fmt.Sprintf("Failed to execute %s", toolName)), nil
327359
}
328360

@@ -346,9 +378,9 @@ func PaginatedToolHandler[T any](manager *kc.Manager, toolName string, apiCall f
346378
})
347379

348380
if err != nil {
349-
handler.trackToolError(toolName, "execution_error")
381+
handler.trackToolError(ctx, toolName, "execution_error")
350382
} else if result != nil && result.IsError {
351-
handler.trackToolError(toolName, "api_error")
383+
handler.trackToolError(ctx, toolName, "api_error")
352384
}
353385
return result, err
354386
}

‎mcp/get_tools.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,7 @@ func (*OrderTradesTool) Tool() mcp.Tool {
207207
func (*OrderTradesTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
208208
handler := NewToolHandler(manager)
209209
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
210+
handler.trackToolCall(ctx, "get_order_trades")
210211
args := request.GetArguments()
211212

212213
// Validate required parameters
@@ -242,6 +243,7 @@ func (*OrderHistoryTool) Tool() mcp.Tool {
242243
func (*OrderHistoryTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
243244
handler := NewToolHandler(manager)
244245
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
246+
handler.trackToolCall(ctx, "get_order_history")
245247
args := request.GetArguments()
246248

247249
// Validate required parameters

‎mcp/market_tools.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ func (*QuotesTool) Tool() mcp.Tool {
2929
func (*QuotesTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
3030
handler := NewToolHandler(manager)
3131
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
32+
handler.trackToolCall(ctx, "get_quotes")
3233
args := request.GetArguments()
3334

3435
// Validate required parameters
@@ -74,6 +75,7 @@ func (*InstrumentsSearchTool) Tool() mcp.Tool {
7475
func (*InstrumentsSearchTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
7576
handler := NewToolHandler(manager)
7677
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
78+
handler.trackToolCall(ctx, "search_instruments")
7779
args := request.GetArguments()
7880

7981
// Validate required parameters
@@ -193,6 +195,7 @@ func (*HistoricalDataTool) Tool() mcp.Tool {
193195
func (*HistoricalDataTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
194196
handler := NewToolHandler(manager)
195197
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
198+
handler.trackToolCall(ctx, "get_historical_data")
196199
args := request.GetArguments()
197200

198201
// Validate required parameters
@@ -256,6 +259,7 @@ func (*LTPTool) Tool() mcp.Tool {
256259
func (*LTPTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
257260
handler := NewToolHandler(manager)
258261
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
262+
handler.trackToolCall(ctx, "get_ltp")
259263
args := request.GetArguments()
260264

261265
// Validate required parameters
@@ -297,6 +301,7 @@ func (*OHLCTool) Tool() mcp.Tool {
297301
func (*OHLCTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
298302
handler := NewToolHandler(manager)
299303
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
304+
handler.trackToolCall(ctx, "get_ohlc")
300305
args := request.GetArguments()
301306

302307
// Validate required parameters

‎mcp/post_tools.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ func (*PlaceOrderTool) Tool() mcp.Tool {
8383
func (*PlaceOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
8484
handler := NewToolHandler(manager)
8585
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
86-
handler.trackToolCall("place_order")
86+
handler.trackToolCall(ctx, "place_order")
8787
args := request.GetArguments()
8888

8989
// Validate required parameters
@@ -165,7 +165,7 @@ func (*ModifyOrderTool) Tool() mcp.Tool {
165165
func (*ModifyOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
166166
handler := NewToolHandler(manager)
167167
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
168-
handler.trackToolCall("modify_order")
168+
handler.trackToolCall(ctx, "modify_order")
169169
args := request.GetArguments()
170170

171171
// Validate required parameters
@@ -218,7 +218,7 @@ func (*CancelOrderTool) Tool() mcp.Tool {
218218
func (*CancelOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
219219
handler := NewToolHandler(manager)
220220
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
221-
handler.trackToolCall("cancel_order")
221+
handler.trackToolCall(ctx, "cancel_order")
222222
args := request.GetArguments()
223223

224224
// Validate required parameters
@@ -310,7 +310,7 @@ func (*PlaceGTTOrderTool) Tool() mcp.Tool {
310310
func (*PlaceGTTOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
311311
handler := NewToolHandler(manager)
312312
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
313-
handler.trackToolCall("place_gtt_order")
313+
handler.trackToolCall(ctx, "place_gtt_order")
314314
args := request.GetArguments()
315315

316316
// Validate required parameters
@@ -383,7 +383,7 @@ func (*DeleteGTTOrderTool) Tool() mcp.Tool {
383383
func (*DeleteGTTOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
384384
handler := NewToolHandler(manager)
385385
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
386-
handler.trackToolCall("delete_gtt_order")
386+
handler.trackToolCall(ctx, "delete_gtt_order")
387387
args := request.GetArguments()
388388

389389
// Validate required parameters
@@ -474,7 +474,7 @@ func (*ModifyGTTOrderTool) Tool() mcp.Tool {
474474
func (*ModifyGTTOrderTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
475475
handler := NewToolHandler(manager)
476476
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
477-
handler.trackToolCall("modify_gtt_order")
477+
handler.trackToolCall(ctx, "modify_gtt_order")
478478
args := request.GetArguments()
479479

480480
// Validate required parameters

0 commit comments

Comments
 (0)