Skip to content

Commit cb81e01

Browse files
rhnvrmrhnvrm
committed
feat: add daily tool usage metrics and AI risk warning
- Add per-day tool usage tracking with date labels for growth analysis - Track tool_calls_{tool_name}_total{date='YYYY-MM-DD'} for all tools - Track tool_errors_{tool_name}_{error_type}_total{date='YYYY-MM-DD'} for errors - Add AI risk warning to login tool: 'AI systems are unpredictable and non-deterministic' - Include comprehensive test suite for daily metrics functionality Enables day-over-day usage comparison and identifies popular tools over time.
1 parent 4113eba commit cb81e01

File tree

5 files changed

+267
-5
lines changed

5 files changed

+267
-5
lines changed

‎app/metrics/daily_metrics_test.go‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
package metrics
2+
3+
import (
4+
"bytes"
5+
"fmt"
6+
"strings"
7+
"testing"
8+
"time"
9+
)
10+
11+
func TestDailyMetrics(t *testing.T) {
12+
m := New(Config{
13+
ServiceName: "test-service",
14+
AutoCleanup: false,
15+
})
16+
17+
today := time.Now().UTC().Format("2006-01-02")
18+
19+
// Test daily increment methods
20+
m.IncrementDaily("tool_calls_quotes")
21+
m.IncrementDaily("tool_calls_quotes")
22+
m.IncrementDailyBy("tool_calls_login", 3)
23+
m.IncrementDaily("tool_errors_quotes_api_error")
24+
25+
// Test regular increment (non-daily)
26+
m.Increment("legacy_counter")
27+
28+
// Check internal storage with date suffix
29+
expectedQuotesKey := fmt.Sprintf("tool_calls_quotes_%s", today)
30+
expectedLoginKey := fmt.Sprintf("tool_calls_login_%s", today)
31+
expectedErrorKey := fmt.Sprintf("tool_errors_quotes_api_error_%s", today)
32+
33+
if got := m.GetCounterValue(expectedQuotesKey); got != 2 {
34+
t.Errorf("Expected tool_calls_quotes_%s = 2, got %d", today, got)
35+
}
36+
37+
if got := m.GetCounterValue(expectedLoginKey); got != 3 {
38+
t.Errorf("Expected tool_calls_login_%s = 3, got %d", today, got)
39+
}
40+
41+
if got := m.GetCounterValue(expectedErrorKey); got != 1 {
42+
t.Errorf("Expected tool_errors_quotes_api_error_%s = 1, got %d", today, got)
43+
}
44+
45+
if got := m.GetCounterValue("legacy_counter"); got != 1 {
46+
t.Errorf("Expected legacy_counter = 1, got %d", got)
47+
}
48+
}
49+
50+
func TestDailyMetricsPrometheusOutput(t *testing.T) {
51+
m := New(Config{
52+
ServiceName: "test-service",
53+
AutoCleanup: false,
54+
})
55+
56+
today := time.Now().UTC().Format("2006-01-02")
57+
58+
// Add some daily metrics
59+
m.IncrementDaily("tool_calls_quotes")
60+
m.IncrementDaily("tool_calls_quotes")
61+
m.IncrementDaily("tool_errors_quotes_api_error")
62+
63+
// Add a regular metric
64+
m.Increment("legacy_counter")
65+
66+
// Generate Prometheus output
67+
var buf bytes.Buffer
68+
m.WritePrometheus(&buf)
69+
output := buf.String()
70+
71+
// Check that daily metrics have date labels
72+
expectedQuotesLine := fmt.Sprintf(`tool_calls_quotes_total{date="%s",service="test-service"} 2`, today)
73+
expectedErrorLine := fmt.Sprintf(`tool_errors_quotes_api_error_total{date="%s",service="test-service"} 1`, today)
74+
expectedLegacyLine := `legacy_counter_total{service="test-service"} 1`
75+
76+
if !strings.Contains(output, expectedQuotesLine) {
77+
t.Errorf("Expected output to contain: %s\nGot: %s", expectedQuotesLine, output)
78+
}
79+
80+
if !strings.Contains(output, expectedErrorLine) {
81+
t.Errorf("Expected output to contain: %s\nGot: %s", expectedErrorLine, output)
82+
}
83+
84+
if !strings.Contains(output, expectedLegacyLine) {
85+
t.Errorf("Expected output to contain: %s\nGot: %s", expectedLegacyLine, output)
86+
}
87+
}
88+
89+
func TestIsDailyMetric(t *testing.T) {
90+
m := New(Config{ServiceName: "test"})
91+
92+
tests := []struct {
93+
key string
94+
expected bool
95+
}{
96+
{"tool_calls_quotes_2025-08-05", true},
97+
{"tool_errors_login_session_error_2025-12-31", true},
98+
{"legacy_counter", false},
99+
{"tool_calls_quotes", false},
100+
{"tool_calls_quotes_20250805", false}, // Wrong date format
101+
{"tool_calls_quotes_2025-8-5", false}, // Wrong date format
102+
{"", false},
103+
{"_2025-08-05", false}, // Empty base name
104+
}
105+
106+
for _, tt := range tests {
107+
t.Run(tt.key, func(t *testing.T) {
108+
if got := m.isDailyMetric(tt.key); got != tt.expected {
109+
t.Errorf("isDailyMetric(%q) = %v, want %v", tt.key, got, tt.expected)
110+
}
111+
})
112+
}
113+
}
114+
115+
func TestParseDailyMetric(t *testing.T) {
116+
m := New(Config{ServiceName: "test"})
117+
118+
tests := []struct {
119+
key string
120+
expectedBase string
121+
expectedDate string
122+
}{
123+
{"tool_calls_quotes_2025-08-05", "tool_calls_quotes", "2025-08-05"},
124+
{"tool_errors_login_session_error_2025-12-31", "tool_errors_login_session_error", "2025-12-31"},
125+
{"legacy_counter", "", ""}, // Not a daily metric
126+
{"tool_calls_quotes_20250805", "", ""}, // Wrong date format
127+
}
128+
129+
for _, tt := range tests {
130+
t.Run(tt.key, func(t *testing.T) {
131+
gotBase, gotDate := m.parseDailyMetric(tt.key)
132+
if gotBase != tt.expectedBase || gotDate != tt.expectedDate {
133+
t.Errorf("parseDailyMetric(%q) = (%q, %q), want (%q, %q)",
134+
tt.key, gotBase, gotDate, tt.expectedBase, tt.expectedDate)
135+
}
136+
})
137+
}
138+
}

‎app/metrics/metrics.go‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,19 @@ func (m *Manager) IncrementBy(key string, n int64) {
9090
atomic.AddInt64(val.(*int64), n)
9191
}
9292

93+
// IncrementDaily atomically increments a daily counter for today
94+
func (m *Manager) IncrementDaily(key string) {
95+
m.IncrementDailyBy(key, 1)
96+
}
97+
98+
// IncrementDailyBy atomically increments a daily counter by n for today
99+
func (m *Manager) IncrementDailyBy(key string, n int64) {
100+
today := time.Now().UTC().Format("2006-01-02")
101+
dailyKey := fmt.Sprintf("%s_%s", key, today)
102+
val, _ := m.counters.LoadOrStore(dailyKey, new(int64))
103+
atomic.AddInt64(val.(*int64), n)
104+
}
105+
93106
// TrackDailyUser tracks a unique user login for today
94107
func (m *Manager) TrackDailyUser(userID string) {
95108
if userID == "" {
@@ -194,6 +207,46 @@ func (m *Manager) Shutdown() {
194207
})
195208
}
196209

210+
// isDailyMetric checks if a metric key has a date suffix (YYYY-MM-DD format)
211+
func (m *Manager) isDailyMetric(key string) bool {
212+
parts := strings.Split(key, "_")
213+
if len(parts) < 2 {
214+
return false
215+
}
216+
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+
223+
// Check if the last part looks like a date (YYYY-MM-DD)
224+
lastPart := parts[len(parts)-1]
225+
if len(lastPart) != 10 || strings.Count(lastPart, "-") != 2 {
226+
return false
227+
}
228+
229+
// Basic validation that it looks like YYYY-MM-DD
230+
dateParts := strings.Split(lastPart, "-")
231+
if len(dateParts) != 3 || len(dateParts[0]) != 4 || len(dateParts[1]) != 2 || len(dateParts[2]) != 2 {
232+
return false
233+
}
234+
235+
return true
236+
}
237+
238+
// parseDailyMetric extracts base name and date from a daily metric key
239+
func (m *Manager) parseDailyMetric(key string) (baseName, date string) {
240+
if !m.isDailyMetric(key) {
241+
return "", ""
242+
}
243+
244+
parts := strings.Split(key, "_")
245+
date = parts[len(parts)-1]
246+
baseName = strings.Join(parts[:len(parts)-1], "_")
247+
return baseName, date
248+
}
249+
197250
// formatMetric formats a single metric in Prometheus format
198251
func (m *Manager) formatMetric(buf *bytes.Buffer, name string, labels map[string]string, value float64) {
199252
if labels == nil {
@@ -215,14 +268,24 @@ func (m *Manager) WritePrometheus(buf *bytes.Buffer) {
215268
now := time.Now().UTC()
216269
today := now.Format("2006-01-02")
217270

218-
// Write counter metrics
271+
// Write counter metrics - separate daily and total counters
219272
m.counters.Range(func(key, val interface{}) bool {
220273
name, ok := key.(string)
221274
if !ok {
222275
return true
223276
}
224277
value := atomic.LoadInt64(val.(*int64))
225-
m.formatMetric(buf, fmt.Sprintf("%s_total", name), nil, float64(value))
278+
279+
// Check if this is a daily metric (has date suffix)
280+
if m.isDailyMetric(name) {
281+
baseName, date := m.parseDailyMetric(name)
282+
if baseName != "" && date != "" {
283+
m.formatMetric(buf, fmt.Sprintf("%s_total", baseName), map[string]string{"date": date}, float64(value))
284+
}
285+
} else {
286+
// Regular total counter without date label
287+
m.formatMetric(buf, fmt.Sprintf("%s_total", name), nil, float64(value))
288+
}
226289
return true
227290
})
228291

‎kc/manager.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,25 @@ func (m *Manager) StopCleanupRoutine() {
431431
m.sessionManager.StopCleanupRoutine()
432432
}
433433

434+
// HasMetrics returns true if metrics manager is available
435+
func (m *Manager) HasMetrics() bool {
436+
return m.metrics != nil
437+
}
438+
439+
// IncrementMetric increments a metric counter by 1
440+
func (m *Manager) IncrementMetric(key string) {
441+
if m.metrics != nil {
442+
m.metrics.Increment(key)
443+
}
444+
}
445+
446+
// IncrementDailyMetric increments a daily metric counter by 1
447+
func (m *Manager) IncrementDailyMetric(key string) {
448+
if m.metrics != nil {
449+
m.metrics.IncrementDaily(key)
450+
}
451+
}
452+
434453
// Shutdown gracefully shuts down the manager and all its components
435454
func (m *Manager) Shutdown() {
436455
m.Logger.Info("Shutting down Kite manager...")

‎mcp/common.go‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,20 @@ func NewToolHandler(manager *kc.Manager) *ToolHandler {
2020
return &ToolHandler{manager: manager}
2121
}
2222

23+
// trackToolCall increments the daily tool usage counter
24+
func (h *ToolHandler) trackToolCall(toolName string) {
25+
if h.manager.HasMetrics() {
26+
h.manager.IncrementDailyMetric(fmt.Sprintf("tool_calls_%s", toolName))
27+
}
28+
}
29+
30+
// trackToolError increments the daily tool error counter with error type
31+
func (h *ToolHandler) trackToolError(toolName, errorType string) {
32+
if h.manager.HasMetrics() {
33+
h.manager.IncrementDailyMetric(fmt.Sprintf("tool_errors_%s_%s", toolName, errorType))
34+
}
35+
}
36+
2337
// WithSession validates session and executes the provided function with a valid Kite session
2438
// This eliminates the TOCTOU race condition by consolidating session validation and usage
2539
func (h *ToolHandler) WithSession(ctx context.Context, toolName string, fn func(*kc.KiteSessionData) (*mcp.CallToolResult, error)) (*mcp.CallToolResult, error) {
@@ -31,11 +45,13 @@ func (h *ToolHandler) WithSession(ctx context.Context, toolName string, fn func(
3145
kiteSession, isNew, err := h.manager.GetOrCreateSession(sessionID)
3246
if err != nil {
3347
h.manager.Logger.Error("Failed to establish session", "tool", toolName, "session_id", sessionID, "error", err)
48+
h.trackToolError(toolName, "session_error")
3449
return mcp.NewToolResultError("Failed to establish a session. Please try again."), nil
3550
}
3651

3752
if isNew {
3853
h.manager.Logger.Info("New session created, login required", "tool", toolName, "session_id", sessionID)
54+
h.trackToolError(toolName, "auth_required")
3955
return mcp.NewToolResultError("Please log in first using the login tool"), nil
4056
}
4157

@@ -283,19 +299,30 @@ func CreatePaginatedResponse(originalData interface{}, paginatedData interface{}
283299
func SimpleToolHandler(manager *kc.Manager, toolName string, apiCall func(*kc.KiteSessionData) (interface{}, error)) server.ToolHandlerFunc {
284300
handler := NewToolHandler(manager)
285301
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
286-
return handler.HandleAPICall(ctx, toolName, apiCall)
302+
// Track the tool call at the handler level
303+
handler.trackToolCall(toolName)
304+
result, err := handler.HandleAPICall(ctx, toolName, apiCall)
305+
if err != nil {
306+
handler.trackToolError(toolName, "execution_error")
307+
} else if result != nil && result.IsError {
308+
handler.trackToolError(toolName, "api_error")
309+
}
310+
return result, err
287311
}
288312
}
289313

290314
// PaginatedToolHandler creates a handler function for endpoints that support pagination
291315
func PaginatedToolHandler[T any](manager *kc.Manager, toolName string, apiCall func(*kc.KiteSessionData) ([]T, error)) server.ToolHandlerFunc {
292316
handler := NewToolHandler(manager)
293317
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
294-
return handler.WithSession(ctx, toolName, func(session *kc.KiteSessionData) (*mcp.CallToolResult, error) {
318+
// Track the tool call at the handler level
319+
handler.trackToolCall(toolName)
320+
result, err := handler.WithSession(ctx, toolName, func(session *kc.KiteSessionData) (*mcp.CallToolResult, error) {
295321
// Get the data
296322
data, err := apiCall(session)
297323
if err != nil {
298324
handler.manager.Logger.Error("API call failed", "tool", toolName, "error", err)
325+
handler.trackToolError(toolName, "api_error")
299326
return mcp.NewToolResultError(fmt.Sprintf("Failed to execute %s", toolName)), nil
300327
}
301328

@@ -317,5 +344,12 @@ func PaginatedToolHandler[T any](manager *kc.Manager, toolName string, apiCall f
317344

318345
return handler.MarshalResponse(responseData, toolName)
319346
})
347+
348+
if err != nil {
349+
handler.trackToolError(toolName, "execution_error")
350+
} else if result != nil && result.IsError {
351+
handler.trackToolError(toolName, "api_error")
352+
}
353+
return result, err
320354
}
321355
}

‎mcp/setup_tools.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ func (*LoginTool) Tool() mcp.Tool {
1919

2020
func (*LoginTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
2121
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
22+
// Track login tool usage
23+
if manager.HasMetrics() {
24+
manager.IncrementDailyMetric("tool_calls_login")
25+
}
26+
2227
// Get MCP client session from context
2328
mcpClientSession := server.ClientSessionFromContext(ctx)
2429

@@ -30,6 +35,9 @@ func (*LoginTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
3035
kiteSession, isNew, err := manager.GetOrCreateSession(mcpSessionID)
3136
if err != nil {
3237
manager.Logger.Error("Failed to get or create Kite session", "session_id", mcpSessionID, "error", err)
38+
if manager.HasMetrics() {
39+
manager.IncrementDailyMetric("tool_errors_login_session_error")
40+
}
3341
return mcp.NewToolResultError("Failed to get or create Kite session"), nil
3442
}
3543

@@ -76,7 +84,7 @@ func (*LoginTool) Handler(manager *kc.Manager) server.ToolHandlerFunc {
7684
Content: []mcp.Content{
7785
mcp.TextContent{
7886
Type: "text",
79-
Text: fmt.Sprintf("Please login to Kite by clicking this link: [Login to Kite](%s)\n\nIf your client supports clickable links, you can render and present it and ask them to click the link above. Otherwise, display the URL and ask them to copy and paste it into their browser: %s\n\nAfter completing the login in your browser, let me know and I'll continue with your request.", url, url),
87+
Text: fmt.Sprintf("IMPORTANT: Please display this warning to the user before proceeding:\n\n⚠️ **WARNING: AI systems are unpredictable and non-deterministic. By continuing, you agree to interact with your Zerodha account via AI at your own risk.**\n\nAfter showing the warning above, provide the user with this login link: [Login to Kite](%s)\n\nIf your client supports clickable links, you can render and present it and ask them to click the link above. Otherwise, display the URL and ask them to copy and paste it into their browser: %s\n\nAfter completing the login in your browser, let me know and I'll continue with your request.", url, url),
8088
},
8189
},
8290
}, nil

0 commit comments

Comments
 (0)