| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "sync/atomic" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | type protocolListStats struct { |
| 10 | durationMs atomic.Int64 |
| 11 | toolCount atomic.Int64 |
| 12 | schemaBytes atomic.Int64 |
| 13 | } |
| 14 | |
| 15 | type protocolMetrics struct { |
| 16 | toolsList atomic.Int64 |
| 17 | toolsCall atomic.Int64 |
| 18 | remote atomic.Int64 |
| 19 | outputSchemaMismatch atomic.Int64 |
| 20 | lists protocolListStats |
| 21 | } |
| 22 | |
| 23 | // ToolsListStats is a snapshot of MCP tools/list observations. |
| 24 | type ToolsListStats struct { |
| 25 | Count, Remote, DurationMs, ToolCount, SchemaBytes int64 |
| 26 | } |
| 27 | |
| 28 | var hostProtocol protocolMetrics |
| 29 | |
| 30 | func (c *Client) observeProtocol(method string, res json.RawMessage, d time.Duration, err error) { |
| 31 | if c == nil { |
| 32 | return |
| 33 | } |
| 34 | switch method { |
| 35 | case "tools/list": |
| 36 | hostProtocol.toolsList.Add(1) |
| 37 | if err == nil { |
| 38 | hostProtocol.remote.Add(1) |
| 39 | n, schemaBytes := countListedTools(res) |
| 40 | hostProtocol.lists.durationMs.Add(d.Milliseconds()) |
| 41 | hostProtocol.lists.toolCount.Add(int64(n)) |
| 42 | hostProtocol.lists.schemaBytes.Add(int64(schemaBytes)) |
| 43 | } |
| 44 | case "tools/call": |
| 45 | if err == nil { |
| 46 | hostProtocol.toolsCall.Add(1) |
| 47 | } |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | func countListedTools(res json.RawMessage) (n, schemaBytes int) { |
| 52 | var out struct { |
| 53 | Tools []struct { |
| 54 | InputSchema json.RawMessage `json:"inputSchema"` |
| 55 | } `json:"tools"` |
| 56 | } |
| 57 | if json.Unmarshal(res, &out) != nil { |
| 58 | return 0, len(res) |
| 59 | } |
| 60 | for _, item := range out.Tools { |
| 61 | schemaBytes += len(item.InputSchema) |
| 62 | } |
| 63 | return len(out.Tools), schemaBytes |
| 64 | } |
| 65 | |
| 66 | func ToolsListCount() int64 { return hostProtocol.toolsList.Load() } |
| 67 | func ToolsCallCount() int64 { return hostProtocol.toolsCall.Load() } |
| 68 | func OutputSchemaMismatchCount() int64 { return hostProtocol.outputSchemaMismatch.Load() } |
| 69 | |
| 70 | // SnapshotToolsListStats returns process-local tools/list counters. |
| 71 | func SnapshotToolsListStats() ToolsListStats { |
| 72 | return ToolsListStats{ |
| 73 | Count: hostProtocol.toolsList.Load(), |
| 74 | Remote: hostProtocol.remote.Load(), |
| 75 | DurationMs: hostProtocol.lists.durationMs.Load(), |
| 76 | ToolCount: hostProtocol.lists.toolCount.Load(), |
| 77 | SchemaBytes: hostProtocol.lists.schemaBytes.Load(), |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | func ResetProtocolMetricsForTest() { |
| 82 | hostProtocol.toolsList.Store(0) |
| 83 | hostProtocol.toolsCall.Store(0) |
| 84 | hostProtocol.remote.Store(0) |
| 85 | hostProtocol.outputSchemaMismatch.Store(0) |
| 86 | hostProtocol.lists.durationMs.Store(0) |
| 87 | hostProtocol.lists.toolCount.Store(0) |
| 88 | hostProtocol.lists.schemaBytes.Store(0) |
| 89 | } |
| 90 |