Coverage for src / infra / telemetry.py: 73%

1""" 

2Telemetry abstraction for agent tracing. 

3 

4Provides a pluggable telemetry system with: 

5- TelemetryProvider protocol for abstraction 

6- TelemetrySpan protocol for span context managers 

7- NullTelemetryProvider for testing and opt-out 

8 

9For production with Braintrust, use BraintrustProvider from braintrust_integration.py: 

10 

11 from src.infra.clients.braintrust_integration import BraintrustProvider 

12 

13Usage: 

14 # For tests or opt-out: 

15 provider = NullTelemetryProvider() 

16 

17 # Use via protocol: 

18 if provider.is_enabled(): 

19 with provider.create_span("task-123", {"agent_id": "agent-1"}): 

20 # Work happens here 

21 pass 

22 provider.flush() 

23""" 

24 

25from __future__ import annotations 

26 

27from typing import TYPE_CHECKING, Any, Protocol, Self 

28 

29if TYPE_CHECKING: 

30 from types import TracebackType 

31 

32 

33class TelemetrySpan(Protocol): 

34 """Protocol for a telemetry span context manager.""" 

35 

36 def __enter__(self) -> Self: 

37 """Enter the span context.""" 

38 ... 

39 

40 def __exit__( 

41 self, 

42 exc_type: type[BaseException] | None, 

43 exc_val: BaseException | None, 

44 exc_tb: TracebackType | None, 

45 ) -> None: 

46 """Exit the span context.""" 

47 ... 

48 

49 def log_input(self, prompt: str) -> None: 

50 """Log the initial user prompt.""" 

51 ... 

52 

53 def log_message(self, message: object) -> None: 

54 """Log a message from the SDK.""" 

55 ... 

56 

57 def set_success(self, success: bool) -> None: 

58 """Mark the execution as successful or failed.""" 

59 ... 

60 

61 def set_error(self, error: str) -> None: 

62 """Record an error message.""" 

63 ... 

64 

65 

66class TelemetryProvider(Protocol): 

67 """Protocol for telemetry providers. 

68 

69 Telemetry providers abstract the underlying tracing system, 

70 allowing tests to use a null implementation and production 

71 code to use Braintrust or other backends. 

72 """ 

73 

74 def is_enabled(self) -> bool: 

75 """Check if telemetry is active and configured.""" 

76 ... 

77 

78 def create_span( 

79 self, name: str, metadata: dict[str, Any] | None = None 

80 ) -> TelemetrySpan: 

81 """Create a span context manager for tracing an operation. 

82 

83 Args: 

84 name: Span name (typically issue_id for agent executions) 

85 metadata: Optional metadata dict (e.g., agent_id, custom fields) 

86 

87 Returns: 

88 A TelemetrySpan context manager for the operation 

89 """ 

90 ... 

91 

92 def flush(self) -> None: 

93 """Flush pending telemetry data.""" 

94 ... 

95 

96 

97class NullSpan: 

98 """No-op span implementation for testing and opt-out. 

99 

100 All methods are no-ops that return immediately. 

101 """ 

102 

103 def __enter__(self) -> Self: 

104 return self 

105 

106 def __exit__( 

107 self, 

108 exc_type: type[BaseException] | None, 

109 exc_val: BaseException | None, 

110 exc_tb: TracebackType | None, 

111 ) -> None: 

112 pass 

113 

114 def log_input(self, prompt: str) -> None: 

115 pass 

116 

117 def log_message(self, message: object) -> None: 

118 pass 

119 

120 def set_success(self, success: bool) -> None: 

121 pass 

122 

123 def set_error(self, error: str) -> None: 

124 pass 

125 

126 

127class NullTelemetryProvider: 

128 """No-op telemetry provider for testing and opt-out. 

129 

130 This provider is completely stateless and has no side effects. 

131 All operations return immediately without doing anything. 

132 

133 Usage: 

134 provider = NullTelemetryProvider() 

135 with provider.create_span("task-123"): 

136 # Work happens here - nothing is traced 

137 pass 

138 """ 

139 

140 def is_enabled(self) -> bool: 

141 """Always returns False - null provider is never 'enabled'.""" 

142 return False 

143 

144 def create_span( 

145 self, name: str, metadata: dict[str, Any] | None = None 

146 ) -> NullSpan: 

147 """Return a no-op span that ignores all operations.""" 

148 return NullSpan() 

149 

150 def flush(self) -> None: 

151 """No-op flush.""" 

152 pass