Coverage for agentos/dashboard/tracker.py: 0%

1""" 

2追踪状态管理器 — 记录 Agent 运行历史、会话、步骤。 

3 

4数据存储在 ~/.agentos/tracker/ 下，以 JSONL 追加写入。 

5""" 

6 

7from __future__ import annotations 

8 

9import json 

10import time 

11from dataclasses import asdict, dataclass, field 

12from pathlib import Path 

13 

14TRACKER_DIR = Path.home() / ".agentos" / "tracker" 

15 

16 

17@dataclass 

18class StepRecord: 

19 """单步执行记录。""" 

20 

21 step_index: int 

22 step_type: str # "thinking" | "tool_call" | "tool_result" | "final_answer" 

23 detail: str # 步骤内容摘要 

24 duration_ms: float 

25 tokens: int = 0 

26 

27 

28@dataclass 

29class SessionRecord: 

30 """单次会话完整记录。""" 

31 

32 session_id: str 

33 task: str 

34 model: str 

35 provider: str 

36 started_at: float = field(default_factory=time.time) 

37 finished_at: float = 0.0 

38 status: str = "running" # "running" | "completed" | "error" | "cancelled" 

39 steps: list[StepRecord] = field(default_factory=list) 

40 total_tokens: int = 0 

41 total_cost_usd: float = 0.0 

42 error: str = "" 

43 

44 

45class Tracker: 

46 """线程安全的追踪记录器（文件级锁）+ 事件发布。 

47 

48 支持订阅者模式：外部可 subscribe 回调，在 add_step/finish_session 时收到实时事件推送。 

49 """ 

50 

51 _instance: Tracker | None = None 

52 

53 def __init__(self): 

54 TRACKER_DIR.mkdir(parents=True, exist_ok=True) 

55 self._sessions_file = TRACKER_DIR / "sessions.jsonl" 

56 self._active: dict[str, SessionRecord] = {} 

57 self._subscribers: list = [] # 回调列表 

58 

59 @classmethod 

60 def get(cls) -> Tracker: 

61 if cls._instance is None: 

62 cls._instance = cls() 

63 return cls._instance 

64 

65 def start_session( 

66 self, session_id: str, task: str, model: str = "", provider: str = "" 

67 ) -> SessionRecord: 

68 rec = SessionRecord(session_id=session_id, task=task, model=model, provider=provider) 

69 self._active[session_id] = rec 

70 return rec 

71 

72 def add_step( 

73 self, 

74 session_id: str, 

75 step_type: str, 

76 detail: str, 

77 duration_ms: float = 0.0, 

78 tokens: int = 0, 

79 ): 

80 rec = self._active.get(session_id) 

81 if rec is None: 

82 return 

83 step = StepRecord( 

84 step_index=len(rec.steps), 

85 step_type=step_type, 

86 detail=detail, 

87 duration_ms=duration_ms, 

88 tokens=tokens, 

89 ) 

90 rec.steps.append(step) 

91 rec.total_tokens += tokens 

92 self._notify( 

93 "step", 

94 { 

95 "session_id": session_id, 

96 "step_type": step_type, 

97 "detail": detail, 

98 "duration_ms": duration_ms, 

99 "tokens": tokens, 

100 }, 

101 ) 

102 

103 def finish_session( 

104 self, session_id: str, status: str = "completed", error: str = "", total_cost: float = 0.0 

105 ): 

106 rec = self._active.pop(session_id, None) 

107 if rec is None: 

108 return 

109 rec.finished_at = time.time() 

110 rec.status = status 

111 rec.error = error 

112 rec.total_cost_usd = total_cost 

113 with open(self._sessions_file, "a") as f: 

114 f.write(json.dumps(asdict(rec), ensure_ascii=False) + "\n") 

115 self._notify("session_done", asdict(rec)) 

116 

117 def subscribe(self, callback): 

118 """订阅实时事件。callback 接收 (event_type: str, data: dict)。""" 

119 self._subscribers.append(callback) 

120 

121 def unsubscribe(self, callback): 

122 """取消订阅。""" 

123 try: 

124 self._subscribers.remove(callback) 

125 except ValueError: 

126 pass 

127 

128 def _notify(self, event_type: str, data: dict): 

129 for cb in self._subscribers: 

130 try: 

131 cb(event_type, data) 

132 except Exception: 

133 pass 

134 

135 def list_sessions(self, limit: int = 50) -> list[dict]: 

136 sessions = [] 

137 if self._sessions_file.exists(): 

138 with open(self._sessions_file) as f: 

139 for line in f: 

140 if line.strip(): 

141 sessions.append(json.loads(line)) 

142 # 倒序，最新的在前 

143 sessions.reverse() 

144 return sessions[:limit] 

145 

146 def get_session(self, session_id: str) -> dict | None: 

147 # 先在 active 中找 

148 rec = self._active.get(session_id) 

149 if rec: 

150 return asdict(rec) 

151 # 再从文件中找 

152 if self._sessions_file.exists(): 

153 with open(self._sessions_file) as f: 

154 for line in f: 

155 if line.strip(): 

156 d = json.loads(line) 

157 if d.get("session_id") == session_id: 

158 return d 

159 return None 

160 

161 def clear(self): 

162 self._active.clear() 

163 if self._sessions_file.exists(): 

164 self._sessions_file.unlink()