1 from __future__ import annotations

2

3 import json

4 import typing as t

5 from dataclasses import dataclass, field

6 from datetime import UTC, datetime

7

8 from mcp_common import MCPServerSettings

9 from mcp_common.cli.health import (

10 RuntimeHealthSnapshot,

11 load_runtime_health,

12 write_runtime_health,

13)

14

15 if t.TYPE_CHECKING:

16 from pathlib import Path

17

18

19@dataclass

20 class RuntimeTelemetrySnapshot:

21 orchestrator_pid: int | None = None

22 started_at: str | None = None

23 updated_at: str | None = None

24 uptime_seconds: float | None = None

25 counters: dict[str, int] = field(default_factory=dict)

26

27 def as_dict(self) -> dict[str, t.Any]:

28 return {

29 "orchestrator_pid": self.orchestrator_pid,

30 "started_at": self.started_at,

31 "updated_at": self.updated_at,

32 "uptime_seconds": self.uptime_seconds,

33 "counters": self.counters.copy(),

34 }

35

36

37@dataclass

38 class RuntimeSnapshotManager:

39 settings: MCPServerSettings

40 started_at: datetime = field(default_factory=lambda: datetime.now(UTC))

41 counters: dict[str, int] = field(default_factory=dict)

42

43 @classmethod

44 def for_server(cls, server_name: str) -> RuntimeSnapshotManager:

45 return cls(settings=MCPServerSettings.load(server_name))

46

47 def record(self, name: str, amount: int = 1) -> None:

48 self.counters[name] = self.counters.get(name, 0) + amount

49

50 def write_health_snapshot(

51 self,

52 pid: int | None,

53 health_state: dict[str, t.Any] | None = None,

54 watchers_running: bool = True,

55 ) -> RuntimeHealthSnapshot:

56 snapshot = load_runtime_health(self.settings.health_snapshot_path())

57 snapshot.orchestrator_pid = pid

58 snapshot.watchers_running = watchers_running

59 if health_state is not None: 59 ↛ 61 line 59 didn't jump to line 61 because the condition on line 59 was always true

60 snapshot.activity_state = {"health": health_state}

61 write_runtime_health(self.settings.health_snapshot_path(), snapshot)

62 return snapshot

63

64 def write_telemetry_snapshot(self, pid: int | None) -> RuntimeTelemetrySnapshot:

65 uptime_seconds = (datetime.now(UTC) - self.started_at).total_seconds()

66 snapshot = RuntimeTelemetrySnapshot(

67 orchestrator_pid=pid,

68 started_at=self.started_at.isoformat(),

69 uptime_seconds=uptime_seconds,

70 counters=self.counters.copy(),

71 )

72 write_runtime_telemetry(self.settings.telemetry_snapshot_path(), snapshot)

73 return snapshot

74

75

76 def load_runtime_telemetry(path: Path) -> RuntimeTelemetrySnapshot:

77 if not path.exists():

78 return RuntimeTelemetrySnapshot()

79

80 try:

81 data = json.loads(path.read_text())

82 except (json.JSONDecodeError, OSError):

83 return RuntimeTelemetrySnapshot()

84

85 if not isinstance(data, dict): 85 ↛ 86 line 85 didn't jump to line 86 because the condition on line 85 was never true

86 return RuntimeTelemetrySnapshot()

87

88 snapshot = RuntimeTelemetrySnapshot()

89 snapshot.orchestrator_pid = data.get("orchestrator_pid")

90 snapshot.started_at = data.get("started_at")

91 snapshot.updated_at = data.get("updated_at")

92 snapshot.uptime_seconds = data.get("uptime_seconds")

93 counters = data.get("counters")

94 snapshot.counters = counters if isinstance(counters, dict) else {}

95 return snapshot

96

97

98 def write_runtime_telemetry(path: Path, snapshot: RuntimeTelemetrySnapshot) -> None:

99 path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)

100 snapshot.updated_at = datetime.now(UTC).isoformat()

101

102 tmp = path.with_suffix(".tmp")

103 try:

104 tmp.write_text(json.dumps(snapshot.as_dict(), indent=2))

105 tmp.chmod(0o600)

106 tmp.replace(path)

107 except OSError:

108 tmp.unlink(missing_ok=True)

109 raise

110

111

112 def update_telemetry_counter(

113 settings: MCPServerSettings,

114 name: str,

115 amount: int = 1,

116 pid: int | None = None,

117) -> RuntimeTelemetrySnapshot:

118 path = settings.telemetry_snapshot_path()

119 snapshot = load_runtime_telemetry(path)

120 now = datetime.now(UTC)

121

122 started_at = _parse_iso_datetime(snapshot.started_at)

123 if started_at is None: 123 ↛ 127 line 123 didn't jump to line 127 because the condition on line 123 was always true

124 started_at = now

125 snapshot.started_at = started_at.isoformat()

126

127 snapshot.orchestrator_pid = pid

128 snapshot.uptime_seconds = (now - started_at).total_seconds()

129 snapshot.counters[name] = snapshot.counters.get(name, 0) + amount

130 write_runtime_telemetry(path, snapshot)

131 return snapshot

132

133

134 async def run_snapshot_loop(

135 manager: RuntimeSnapshotManager,

136 pid: int | None,

137 interval_seconds: float,

138) -> None:

139 while True:

140 manager.record("snapshot_updates")

141 manager.write_health_snapshot(pid=pid)

142 manager.write_telemetry_snapshot(pid=pid)

143 await _sleep(interval_seconds)

144

145

146 async def _sleep(interval_seconds: float) -> None:

147 import asyncio

148

149 await asyncio.sleep(interval_seconds)

150

151

152 def _parse_iso_datetime(value: str | None) -> datetime | None:

153 if value is None: 153 ↛ 155 line 153 didn't jump to line 155 because the condition on line 153 was always true

154 return None

155 try:

156 return datetime.fromisoformat(value)

157 except (TypeError, ValueError):

158 return None

159

160

161__all__ = [

162 "RuntimeSnapshotManager",

163 "RuntimeTelemetrySnapshot",

164 "load_runtime_telemetry",

165 "run_snapshot_loop",

166 "update_telemetry_counter",

167 "write_runtime_telemetry",

168]