Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
3.9 kB · 135 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136import { isAgentSessionMessage } from "../../core/agent-messages.js";import type { AgentConnectionSessionEvent } from "../agent-connection/index.js";
export type AgentActivity = "waiting" | "thinking" | "writing" | "writing-code" | "executing";
export interface AgentActivityStatus { activity: AgentActivity; direction: "down" | "up"; tokens: number;}
export const AGENT_ACTIVITY_LABELS: Record<AgentActivity, string> = { waiting: "Waiting", thinking: "Thinking", writing: "Writing", "writing-code": "Writing code", executing: "Executing",};
const CHARS_PER_TOKEN_ESTIMATE = 4;
export class AgentActivityTracker { private activity: AgentActivity = "waiting"; private completedTokens = 0; private streamingUsageTokens = 0; private streamingChars = 0; private runningToolCount = 0; // Providers like Anthropic only report usage at the start and end of a message, so the // live count leans on the character estimate in between. Keeping the reported value // monotonic prevents it from dipping when authoritative usage arrives at message end. private reportedTokens = 0;
handleEvent(event: AgentConnectionSessionEvent): void { switch (event.type) { case "agent_start": this.activity = "waiting"; this.runningToolCount = 0; break;
case "message_start": if (event.message.role === "user" || isAgentSessionMessage(event.message)) { this.reset(); } else if (event.message.role === "assistant") { this.activity = "waiting"; this.streamingUsageTokens = 0; this.streamingChars = 0; } break;
case "message_update": { if (event.message.role !== "assistant") break; const streamEvent = event.assistantMessageEvent; switch (streamEvent.type) { case "thinking_start": case "thinking_delta": this.activity = "thinking"; break; case "text_start": case "text_delta": this.activity = "writing"; break; case "toolcall_start": case "toolcall_delta": this.activity = "writing-code"; break; default: break; } if ("delta" in streamEvent) { this.streamingChars += streamEvent.delta.length; } this.streamingUsageTokens = event.message.usage.output; break; }
case "message_end": if (event.message.role !== "assistant") break; this.completedTokens += event.message.usage.output > 0 ? event.message.usage.output : this.estimatedStreamingTokens(); this.streamingUsageTokens = 0; this.streamingChars = 0; this.activity = "waiting"; break;
case "tool_execution_start": this.runningToolCount++; this.activity = "executing"; break;
case "tool_execution_end": this.runningToolCount = Math.max(0, this.runningToolCount - 1); if (this.runningToolCount === 0) { this.activity = "waiting"; } break;
default: break; } this.reportedTokens = Math.max(this.reportedTokens, this.currentTokens()); }
getStatus(): AgentActivityStatus { return { activity: this.activity, direction: this.activity === "waiting" || this.activity === "executing" ? "up" : "down", tokens: this.reportedTokens, }; }
private currentTokens(): number { return this.completedTokens + Math.max(this.streamingUsageTokens, this.estimatedStreamingTokens()); }
reset(): void { this.activity = "waiting"; this.completedTokens = 0; this.streamingUsageTokens = 0; this.streamingChars = 0; this.runningToolCount = 0; this.reportedTokens = 0; }
private estimatedStreamingTokens(): number { return Math.round(this.streamingChars / CHARS_PER_TOKEN_ESTIMATE); }}
export function formatTokenCount(count: number): string { if (count < 1000) return count.toString(); if (count < 10000) return `${(count / 1000).toFixed(1)}k`; if (count < 1000000) return `${Math.round(count / 1000)}k`; if (count < 10000000) return `${(count / 1000000).toFixed(1)}M`; return `${Math.round(count / 1000000)}M`;}