Skip to content
Source
Sign inSign up
Open History to see the last commit

usage.ts

185 lines6.7 KB

1import * as Effect from "effect/Effect";
2import * as Option from "effect/Option";
3import * as Path from "effect/Path";
4import * as Schema from "effect/Schema";
5​
6import { expandHomePath } from "@t3tools/provider-core/server/pathExpansion";
7import {
8 parseTimestampMs,
9 tokenCount,
10 totalTokens,
11 type ProviderUsageReader,
12 type TranscriptUsageFormat,
13 type UsageRecord,
14} from "@t3tools/provider-core/server/usage";
15import * as HostProcess from "@t3tools/shared/HostProcess";
16import type { PiSettings } from "../settings.ts";
17import { PI_ARGUMENTS_WITH_VALUES, resolvePiLaunchArgs } from "./mcpInjection.ts";
18​
19const PiScanState = Schema.Struct({ sessionId: Schema.mutableKey(Schema.String) });
20const Usage = Schema.Struct({
21 input: Schema.optional(Schema.Unknown),
22 output: Schema.optional(Schema.Unknown),
23 cacheRead: Schema.optional(Schema.Unknown),
24 cacheWrite: Schema.optional(Schema.Unknown),
25 reasoning: Schema.optional(Schema.Unknown),
26 cost: Schema.optional(Schema.Struct({ total: Schema.Unknown })),
27});
28const Identity = { id: Schema.NonEmptyString, timestamp: Schema.String };
29const Model = { provider: Schema.NonEmptyString, model: Schema.NonEmptyString };
30const Entry = Schema.Union([
31 Schema.Struct({ type: Schema.Literal("session"), id: Schema.NonEmptyString }),
32 Schema.Struct({
33 ...Identity,
34 type: Schema.Literal("message"),
35 message: Schema.Union([
36 Schema.Struct({
37 ...Model,
38 role: Schema.Literal("assistant"),
39 responseModel: Schema.optional(Schema.NonEmptyString),
40 usage: Usage,
41 }),
42 Schema.Struct({ role: Schema.Literal("toolResult"), usage: Usage }),
43 ]),
44 }),
45 Schema.Struct({ ...Identity, ...Model, type: Schema.Literal("usage"), usage: Usage }),
46 Schema.Struct({
47 ...Identity,
48 type: Schema.Literals(["compaction", "branch_summary"]),
49 usage: Usage,
50 }),
51]);
52const decodeEntry = Schema.decodeUnknownOption(Entry);
53​
54function parseRecord(parsed: unknown, state: typeof PiScanState.Type): readonly UsageRecord[] {
55 const decoded = decodeEntry(parsed);
56 if (Option.isNone(decoded)) return [];
57 const entry = decoded.value;
58 if (entry.type === "session") {
59 state.sessionId = entry.id;
60 return [];
61 }
62 const timestampMs = parseTimestampMs(entry.timestamp);
63 if (timestampMs === null) return [];
64 const source = entry.type === "message" ? entry.message : entry;
65 const usage = source.usage;
66 // Auxiliary tool/summary records lack model attribution; match Pi's breakdown.
67 const rateModel =
68 "model" in source
69 ? "responseModel" in source
70 ? (source.responseModel ?? source.model)
71 : source.model
72 : undefined;
73 const model = "provider" in source ? `${source.provider}/${rateModel}` : "Tools/summaries";
74 const outputTokens = tokenCount(usage.output);
75 const totals = {
76 // Pi's input/cache counters are already disjoint.
77 uncachedInputTokens: tokenCount(usage.input),
78 cachedInputTokens: tokenCount(usage.cacheRead),
79 cacheCreationTokens: tokenCount(usage.cacheWrite),
80 outputTokens,
81 reasoningTokens: Math.min(outputTokens, tokenCount(usage.reasoning)),
82 };
83 const cost = usage.cost?.total;
84 // Omitted Pi rates default to zero; use shared estimates without positive transcript cost.
85 const reportedCostUsd =
86 typeof cost === "number" && Number.isFinite(cost) && cost > 0 ? cost : null;
87 if (totalTokens(totals) === 0 && reportedCostUsd === null) return [];
88 return [
89 {
90 provider: "pi",
91 timestampMs,
92 model,
93 ...(rateModel === undefined ? {} : { rateModel }),
94 sessionId: state.sessionId,
95 totals,
96 reportedCostUsd,
97 speed: "standard",
98 // Forks preserve IDs/timestamps; short random IDs need timestamp/model to avoid cross-session collisions.
99 dedupeKey: JSON.stringify(["pi", entry.id, entry.timestamp, model]),
100 },
101 ];
103​
104export const piUsageFormat: TranscriptUsageFormat<typeof PiScanState.Type> = {
105 selectFields: {
106 type: true,
107 id: true,
108 timestamp: true,
109 provider: true,
110 model: true,
111 usage: true,
112 message: { role: true, provider: true, model: true, responseModel: true, usage: true },
113 },
114 usageMarkers: ['"usage"', '"session"'],
115 parseLine: (line, state) => {
116 let parsed: unknown;
117 try {
118 parsed = JSON.parse(line);
119 } catch {
120 return [];
121 }
122 return parseRecord(parsed, state);
123 },
124 parseProjected: parseRecord,
125 state: { initial: () => ({ sessionId: "" }), schema: PiScanState },
126};
127​
128export const piUsageReader: ProviderUsageReader<PiSettings, Path.Path> = {
129 kind: "transcripts",
130 provider: "pi",
131 format: piUsageFormat,
132 directories: Effect.fn("piUsageReader.directories")(function* ({ config, environment }) {
133 const path = yield* Path.Path;
134 const home = yield* HostProcess.HomeDirectory;
135 const platform = yield* HostProcess.Platform;
136 const launchHome = piHome(environment, home, platform);
137 const sessionDirArg = config === undefined ? undefined : piLaunchSessionDir(config.launchArgs);
138 const launchDir =
139 sessionDirArg === undefined ? undefined : expandHomePath(sessionDirArg, launchHome);
140 // --session-dir moves only T3-launched sessions; terminal sessions stay in Pi's own directory.
141 return [
142 ...(launchDir !== undefined && path.isAbsolute(launchDir)
143 ? [{ dir: path.resolve(launchDir) }]
144 : []),
145 { dir: resolvePiSessionsDir(environment, launchHome, path) },
146 ];
147 }),
148};
149​
150function piLaunchSessionDir(launchArgs: string): string | undefined {
151 const resolved = resolvePiLaunchArgs(launchArgs);
152 if (!resolved.ok) return undefined;
153 let sessionDir: string | undefined;
154 for (let index = 0; index < resolved.args.length; index += 1) {
155 const arg = resolved.args[index];
156 if (arg === undefined || !PI_ARGUMENTS_WITH_VALUES.has(arg)) continue;
157 index += 1;
158 if (arg === "--session-dir") sessionDir = resolved.args[index] || undefined;
159 }
160 return sessionDir;
162​
163/** The home Pi resolves `~` against: the instance's HOME (USERPROFILE on Windows), else the host's. */
164export function piHome(
165 environment: NodeJS.ProcessEnv,
166 hostHome: string,
167 platform: NodeJS.Platform,
168): string {
169 return (platform === "win32" ? environment.USERPROFILE : environment.HOME) || hostHome;
171​
172/** Where Pi stores session JSONL files for a provider environment, matching Pi's own lookup. */
173export function resolvePiSessionsDir(
174 environment: NodeJS.ProcessEnv,
175 home: string,
176 path: Path.Path,
177): string {
178 const sessionDir = environment.PI_CODING_AGENT_SESSION_DIR?.trim();
179 if (sessionDir) return path.resolve(expandHomePath(sessionDir, home));
180 const agentDir = environment.PI_CODING_AGENT_DIR?.trim();
181 return path.resolve(
182 agentDir ? expandHomePath(agentDir, home) : path.join(home, ".pi", "agent"),
183 "sessions",
184 );