@nmzpy/pi-ember-stack 0.1.5 → 0.2.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,87 @@
1
+ import path from "node:path";
2
+
3
+ export function normalizePathConstraint(
4
+ pathConstraint: string,
5
+ cwd = process.cwd(),
6
+ ): string | null {
7
+ let trimmed = pathConstraint.trim();
8
+ if (!trimmed) return trimmed;
9
+
10
+ if (path.isAbsolute(trimmed)) {
11
+ const relative = path.relative(cwd, trimmed).replaceAll(path.sep, "/");
12
+ if (relative === "") return null;
13
+ if (relative.startsWith("../") || relative === ".." || path.isAbsolute(relative)) {
14
+ throw new Error(
15
+ `Path constraint must be relative to the workspace: ${pathConstraint}`,
16
+ );
17
+ }
18
+ trimmed = relative;
19
+ }
20
+
21
+ if (trimmed === "." || trimmed === "./") return null;
22
+ // Strip a leading `./` so `./**/*.rs` and `**/*.rs` behave identically.
23
+ if (trimmed.startsWith("./")) trimmed = trimmed.slice(2);
24
+
25
+ // FFF's glob matcher can treat a hidden directory root glob such as
26
+ // `.agents/**` as empty, while the tool contract says this means "inside
27
+ // this directory". Collapse simple trailing recursive directory globs to the
28
+ // directory-prefix constraint understood by the parser. Keep real file globs
29
+ // such as `src/**/*.ts` unchanged.
30
+ const recursiveDir = trimmed.match(/^(.*)\/\*\*(?:\/\*)?$/);
31
+ if (recursiveDir) {
32
+ const dir = recursiveDir[1];
33
+ if (dir && !/[*?[{]/.test(dir)) return `${dir}/`;
34
+ }
35
+
36
+ // Already signals path-constraint syntax to the parser.
37
+ if (trimmed.startsWith("/") || trimmed.endsWith("/")) return trimmed;
38
+ // Globs (`*.ts`, `src/**/*.cc`, `{src,lib}`) are handled by the parser.
39
+ if (/[*?[{]/.test(trimmed)) return trimmed;
40
+ // Filename with extension (`main.rs`, `config.json`) → FilePath constraint.
41
+ const lastSegment = trimmed.split("/").pop() ?? "";
42
+ if (/\.[a-zA-Z][a-zA-Z0-9]{0,9}$/.test(lastSegment)) return trimmed;
43
+ // Bare directory prefix → append `/` so the parser sees a PathSegment.
44
+ return `${trimmed}/`;
45
+ }
46
+
47
+ // Exclusions are emitted as `!<constraint>` tokens, which the Rust parser
48
+ // understands (crates/fff-query-parser/src/parser.rs). We normalize each one
49
+ // the same way as the include path so bare dirs become PathSegment excludes.
50
+ // Tolerate callers passing already-negated forms like `!src/` by stripping
51
+ // the leading `!` before normalizing so we never double-negate (`!!src/`).
52
+ export function normalizeExcludes(
53
+ exclude: string | string[] | undefined,
54
+ cwd = process.cwd(),
55
+ ): string[] {
56
+ if (!exclude) return [];
57
+ const list = Array.isArray(exclude) ? exclude : [exclude];
58
+ const out: string[] = [];
59
+ for (const raw of list) {
60
+ const parts = raw
61
+ .split(/[,\s]+/)
62
+ .map((s) => s.trim())
63
+ .filter(Boolean);
64
+ for (const p of parts) {
65
+ const stripped = p.startsWith("!") ? p.slice(1) : p;
66
+ const normalized = normalizePathConstraint(stripped, cwd);
67
+ if (normalized) out.push(`!${normalized}`);
68
+ }
69
+ }
70
+ return out;
71
+ }
72
+
73
+ export function buildQuery(
74
+ path: string | undefined,
75
+ pattern: string,
76
+ exclude?: string | string[],
77
+ cwd = process.cwd(),
78
+ ): string {
79
+ const parts: string[] = [];
80
+ if (path) {
81
+ const pathConstraint = normalizePathConstraint(path, cwd);
82
+ if (pathConstraint) parts.push(pathConstraint);
83
+ }
84
+ parts.push(...normalizeExcludes(exclude, cwd));
85
+ parts.push(pattern);
86
+ return parts.join(" ");
87
+ }
@@ -0,0 +1,66 @@
1
+ import { describe, expect, test } from "bun:test";
2
+ import { normalizePathConstraint, normalizeExcludes, buildQuery } from "../query.ts";
3
+
4
+ describe("normalizePathConstraint", () => {
5
+ test("empty string returns empty", () => {
6
+ expect(normalizePathConstraint("")).toBe("");
7
+ });
8
+
9
+ test("dot returns null", () => {
10
+ expect(normalizePathConstraint(".")).toBeNull();
11
+ expect(normalizePathConstraint("./")).toBeNull();
12
+ });
13
+
14
+ test("bare directory gets trailing slash", () => {
15
+ expect(normalizePathConstraint("src")).toBe("src/");
16
+ });
17
+
18
+ test("glob is preserved", () => {
19
+ expect(normalizePathConstraint("*.ts")).toBe("*.ts");
20
+ expect(normalizePathConstraint("src/**/*.cc")).toBe("src/**/*.cc");
21
+ });
22
+
23
+ test("filename with extension is preserved", () => {
24
+ expect(normalizePathConstraint("main.rs")).toBe("main.rs");
25
+ });
26
+
27
+ test("recursive dir glob collapses to prefix", () => {
28
+ expect(normalizePathConstraint("src/**")).toBe("src/");
29
+ expect(normalizePathConstraint("src/**/*")).toBe("src/");
30
+ });
31
+
32
+ test("leading ./ is stripped", () => {
33
+ expect(normalizePathConstraint("./src")).toBe("src/");
34
+ });
35
+ });
36
+
37
+ describe("normalizeExcludes", () => {
38
+ test("undefined returns empty array", () => {
39
+ expect(normalizeExcludes(undefined)).toEqual([]);
40
+ });
41
+
42
+ test("comma-separated string is split", () => {
43
+ expect(normalizeExcludes("test/, *.min.js")).toEqual(["!test/", "!*.min.js"]);
44
+ });
45
+
46
+ test("leading ! is tolerated", () => {
47
+ expect(normalizeExcludes("!src/")).toEqual(["!src/"]);
48
+ });
49
+
50
+ test("array input works", () => {
51
+ expect(normalizeExcludes(["test/", "vendor/"])).toEqual(["!test/", "!vendor/"]);
52
+ });
53
+ });
54
+
55
+ describe("buildQuery", () => {
56
+ test("assembles path + excludes + pattern", () => {
57
+ const q = buildQuery("src/", "MyClass", "test/");
58
+ expect(q).toContain("src/");
59
+ expect(q).toContain("!test/");
60
+ expect(q).toContain("MyClass");
61
+ });
62
+
63
+ test("no path or excludes yields just pattern", () => {
64
+ expect(buildQuery(undefined, "foo")).toBe("foo");
65
+ });
66
+ });
@@ -0,0 +1,376 @@
1
+ import { describe, expect, mock, test } from "bun:test";
2
+ import { CompactRenderer, formatCallBody, DISCOVERY_TOOLS } from "../../pi-compact-tools/renderer.ts";
3
+
4
+ function makeTheme() {
5
+ const fg = mock((tag: string, text: string) => `[${tag}:${text}]`);
6
+ return {
7
+ fg,
8
+ bold: mock((s: string) => `*${s}*`),
9
+ };
10
+ }
11
+
12
+ function makeContext(id: string, state: Record<string, any> = {}) {
13
+ return {
14
+ args: {},
15
+ toolCallId: id,
16
+ invalidate: mock(() => {}),
17
+ state,
18
+ };
19
+ }
20
+
21
+ describe("CompactRenderer", () => {
22
+ test("beginTurn does not reset grouping for same-key calls", () => {
23
+ const r = new CompactRenderer();
24
+ const theme = makeTheme() as any;
25
+ const ctx1 = makeContext("a");
26
+ r.renderCall("read", { file_path: "foo.ts" }, theme, ctx1 as any);
27
+ r.beginTurn();
28
+ const ctx2 = makeContext("b");
29
+ r.renderCall("read", { file_path: "bar.ts" }, theme, ctx2 as any);
30
+ // Same group key (__discovery__) should still group across turns
31
+ const record = (r as any).calls.get("b");
32
+ expect(record.group).toBeDefined();
33
+ });
34
+
35
+ test("consecutive discovery tools group together", () => {
36
+ const r = new CompactRenderer();
37
+ const theme = makeTheme() as any;
38
+ const ctx1 = makeContext("a");
39
+ const ctx2 = makeContext("b");
40
+ r.renderCall("read", { file_path: "a.ts" }, theme, ctx1 as any);
41
+ r.renderCall("grep", { pattern: "foo" }, theme, ctx2 as any);
42
+ const recA = (r as any).calls.get("a");
43
+ const recB = (r as any).calls.get("b");
44
+ expect(recA.group).toBeDefined();
45
+ expect(recB.group).toBe(recA.group);
46
+ expect(recA.group.records.length).toBe(2);
47
+ });
48
+
49
+ test("non-discovery tools do not group", () => {
50
+ const r = new CompactRenderer();
51
+ const theme = makeTheme() as any;
52
+ r.renderCall("bash", { command: "ls" }, theme, makeContext("a") as any);
53
+ r.renderCall("edit", { file_path: "x.ts" }, theme, makeContext("b") as any);
54
+ const recA = (r as any).calls.get("a");
55
+ const recB = (r as any).calls.get("b");
56
+ expect(recA.group).toBeUndefined();
57
+ expect(recB.group).toBeUndefined();
58
+ });
59
+
60
+ test("renderResult returns empty for bash on success", () => {
61
+ const r = new CompactRenderer();
62
+ const theme = makeTheme() as any;
63
+ r.renderCall("bash", { command: "echo hi" }, theme, makeContext("a") as any);
64
+ const comp = r.renderResult(
65
+ "bash",
66
+ { command: "echo hi" },
67
+ { content: [{ type: "text", text: "hi" }] },
68
+ { isPartial: false },
69
+ theme,
70
+ { ...makeContext("a"), isError: false } as any,
71
+ );
72
+ // Compact: the result component itself is empty, but the call row is updated.
73
+ expect((comp as any).text).toBe("");
74
+ });
75
+
76
+ test("bash success renders last output line on the call row", () => {
77
+ const r = new CompactRenderer();
78
+ const theme = makeTheme() as any;
79
+ const state: Record<string, any> = {};
80
+ const call = r.renderCall("bash", { command: "echo hi" }, theme, makeContext("a", state) as any);
81
+ r.renderResult(
82
+ "bash",
83
+ { command: "echo hi" },
84
+ { content: [{ type: "text", text: "first\nsecond\nOK" }] },
85
+ { isPartial: false },
86
+ theme,
87
+ { ...makeContext("a", state), isError: false } as any,
88
+ );
89
+ expect((call as any).text).toContain("Run");
90
+ expect((call as any).text).toContain("OK");
91
+ expect((call as any).text).not.toContain("first");
92
+ expect((call as any).text).not.toContain("second");
93
+ });
94
+
95
+ test("bash output skips trailing blank lines", () => {
96
+ const r = new CompactRenderer();
97
+ const theme = makeTheme() as any;
98
+ const state: Record<string, any> = {};
99
+ const call = r.renderCall("bash", { command: "echo hi" }, theme, makeContext("a", state) as any);
100
+ r.renderResult(
101
+ "bash",
102
+ { command: "echo hi" },
103
+ { content: [{ type: "text", text: "OK\n\n\n" }] },
104
+ { isPartial: false },
105
+ theme,
106
+ { ...makeContext("a", state), isError: false } as any,
107
+ );
108
+ expect((call as any).text).toContain("OK");
109
+ expect((call as any).text).not.toContain("\n\n");
110
+ });
111
+
112
+ test("bash empty output does not add a result line", () => {
113
+ const r = new CompactRenderer();
114
+ const theme = makeTheme() as any;
115
+ const state: Record<string, any> = {};
116
+ const call = r.renderCall("bash", { command: ":" }, theme, makeContext("a", state) as any);
117
+ r.renderResult(
118
+ "bash",
119
+ { command: ":" },
120
+ { content: [{ type: "text", text: "" }] },
121
+ { isPartial: false },
122
+ theme,
123
+ { ...makeContext("a", state), isError: false } as any,
124
+ );
125
+ expect((call as any).text).toContain("Run");
126
+ expect((call as any).text).not.toContain(" ");
127
+ });
128
+
129
+ test("bash partial result does not render output yet", () => {
130
+ const r = new CompactRenderer();
131
+ const theme = makeTheme() as any;
132
+ const state: Record<string, any> = {};
133
+ const call = r.renderCall("bash", { command: "sleep 1" }, theme, makeContext("a", state) as any);
134
+ r.renderResult(
135
+ "bash",
136
+ { command: "sleep 1" },
137
+ { content: [{ type: "text", text: "intermediate" }] },
138
+ { isPartial: true },
139
+ theme,
140
+ { ...makeContext("a", state), isError: false } as any,
141
+ );
142
+ expect((call as any).text).toContain("Run");
143
+ expect((call as any).text).not.toContain("intermediate");
144
+ });
145
+
146
+ test("renderResult shows error text on failure", () => {
147
+ const r = new CompactRenderer();
148
+ const theme = makeTheme() as any;
149
+ r.renderCall("read", { file_path: "bad.ts" }, theme, makeContext("a") as any);
150
+ const comp = r.renderResult(
151
+ "read",
152
+ { file_path: "bad.ts" },
153
+ { content: [{ type: "text", text: "Error: not found" }] },
154
+ { isPartial: false },
155
+ theme,
156
+ { ...makeContext("a"), isError: true } as any,
157
+ );
158
+ expect((comp as any).text).toContain("Error: not found");
159
+ });
160
+
161
+ test("read never dumps file content on success", () => {
162
+ const r = new CompactRenderer();
163
+ const theme = makeTheme() as any;
164
+ r.renderCall("read", { file_path: "big.ts" }, theme, makeContext("a") as any);
165
+ const fileContent = "line1\nline2\nline3\n".repeat(100);
166
+ const comp = r.renderResult(
167
+ "read",
168
+ { file_path: "big.ts" },
169
+ { content: [{ type: "text", text: fileContent }] },
170
+ { isPartial: false },
171
+ theme,
172
+ { ...makeContext("a"), isError: false } as any,
173
+ );
174
+ // Result must be empty — the call row shows the compact summary.
175
+ expect((comp as any).text).toBe("");
176
+ expect((comp as any).text).not.toContain("line1");
177
+ });
178
+
179
+ test("bash calls with same cd dir group across turns", () => {
180
+ const r = new CompactRenderer();
181
+ const theme = makeTheme() as any;
182
+ r.renderCall("bash", { command: "cd src && ls" }, theme, makeContext("a") as any);
183
+ r.beginTurn();
184
+ r.renderCall("bash", { command: "cd src && pwd" }, theme, makeContext("b") as any);
185
+ const recA = (r as any).calls.get("a");
186
+ const recB = (r as any).calls.get("b");
187
+ expect(recA.group).toBeDefined();
188
+ expect(recB.group).toBe(recA.group);
189
+ });
190
+
191
+ test("DISCOVERY_TOOLS includes grep and find", () => {
192
+ expect(DISCOVERY_TOOLS.has("grep")).toBe(true);
193
+ expect(DISCOVERY_TOOLS.has("find")).toBe(true);
194
+ expect(DISCOVERY_TOOLS.has("read")).toBe(true);
195
+ expect(DISCOVERY_TOOLS.has("ls")).toBe(true);
196
+ });
197
+
198
+ test("formatCallBody handles grep", () => {
199
+ const theme = makeTheme() as any;
200
+ const result = formatCallBody("grep", { pattern: "MyClass", path: "src/" }, theme);
201
+ expect(result).toContain("Search");
202
+ expect(result).toContain("MyClass");
203
+ expect(result).toContain("src/");
204
+ });
205
+
206
+ test("formatCallBody handles find", () => {
207
+ const theme = makeTheme() as any;
208
+ const result = formatCallBody("find", { pattern: "main", path: "." }, theme);
209
+ expect(result).toContain("Find");
210
+ expect(result).toContain("main");
211
+ });
212
+
213
+ test("renderOwner moves group to latest call", () => {
214
+ const r = new CompactRenderer();
215
+ const theme = makeTheme() as any;
216
+ const ctx1 = makeContext("a");
217
+ const ctx2 = makeContext("b");
218
+ r.renderCall("read", { file_path: "a.ts" }, theme, ctx1 as any);
219
+ r.renderCall("read", { file_path: "b.ts" }, theme, ctx2 as any);
220
+ const recA = (r as any).calls.get("a");
221
+ const recB = (r as any).calls.get("b");
222
+ // After B joins, B should be the renderOwner
223
+ expect(recA.group.renderOwner).toBe(recB);
224
+ expect(recA.group.renderOwner).not.toBe(recA);
225
+ });
226
+
227
+ test("non-owner grouped call renders empty", () => {
228
+ const r = new CompactRenderer();
229
+ const theme = makeTheme() as any;
230
+ const ctx1 = makeContext("a");
231
+ const ctx2 = makeContext("b");
232
+ r.renderCall("read", { file_path: "a.ts" }, theme, ctx1 as any);
233
+ r.renderCall("read", { file_path: "b.ts" }, theme, ctx2 as any);
234
+ // After B joins and becomes owner, re-rendering A should return empty
235
+ const compA2 = r.renderCall("read", { file_path: "a.ts" }, theme, ctx1 as any);
236
+ expect((compA2 as any).text).toBe("");
237
+ });
238
+
239
+ test("expanded bash shows full output", () => {
240
+ const r = new CompactRenderer();
241
+ const theme = makeTheme() as any;
242
+ const state: Record<string, any> = {};
243
+ r.renderCall("bash", { command: "echo hi" }, theme, makeContext("a", state) as any);
244
+ const comp = r.renderResult(
245
+ "bash",
246
+ { command: "echo hi" },
247
+ { content: [{ type: "text", text: "line1\nline2\nline3" }] },
248
+ { isPartial: false, expanded: true },
249
+ theme,
250
+ { ...makeContext("a", state), isError: false } as any,
251
+ );
252
+ expect((comp as any).text).toContain("line1");
253
+ expect((comp as any).text).toContain("line2");
254
+ expect((comp as any).text).toContain("line3");
255
+ });
256
+
257
+ test("collapsed bash shows only last line", () => {
258
+ const r = new CompactRenderer();
259
+ const theme = makeTheme() as any;
260
+ const state: Record<string, any> = {};
261
+ const call = r.renderCall("bash", { command: "echo hi" }, theme, makeContext("a", state) as any);
262
+ r.renderResult(
263
+ "bash",
264
+ { command: "echo hi" },
265
+ { content: [{ type: "text", text: "line1\nline2\nline3" }] },
266
+ { isPartial: false, expanded: false },
267
+ theme,
268
+ { ...makeContext("a", state), isError: false } as any,
269
+ );
270
+ expect((call as any).text).toContain("line3");
271
+ expect((call as any).text).not.toContain("line1");
272
+ expect((call as any).text).not.toContain("line2");
273
+ });
274
+
275
+ test("expanded read shows full file content", () => {
276
+ const r = new CompactRenderer();
277
+ const theme = makeTheme() as any;
278
+ r.renderCall("read", { file_path: "big.ts" }, theme, makeContext("a") as any);
279
+ const fileContent = "line1\nline2\nline3\nline4\nline5";
280
+ const comp = r.renderResult(
281
+ "read",
282
+ { file_path: "big.ts" },
283
+ { content: [{ type: "text", text: fileContent }] },
284
+ { isPartial: false, expanded: true },
285
+ theme,
286
+ { ...makeContext("a"), isError: false } as any,
287
+ );
288
+ expect((comp as any).text).toContain("line1");
289
+ expect((comp as any).text).toContain("line2");
290
+ expect((comp as any).text).toContain("line5");
291
+ });
292
+
293
+ test("collapsed read hides file content", () => {
294
+ const r = new CompactRenderer();
295
+ const theme = makeTheme() as any;
296
+ r.renderCall("read", { file_path: "big.ts" }, theme, makeContext("a") as any);
297
+ const fileContent = "line1\nline2\nline3\nline4\nline5";
298
+ const comp = r.renderResult(
299
+ "read",
300
+ { file_path: "big.ts" },
301
+ { content: [{ type: "text", text: fileContent }] },
302
+ { isPartial: false, expanded: false },
303
+ theme,
304
+ { ...makeContext("a"), isError: false } as any,
305
+ );
306
+ expect((comp as any).text).toBe("");
307
+ });
308
+
309
+ test("expanded grep shows all match lines", () => {
310
+ const r = new CompactRenderer();
311
+ const theme = makeTheme() as any;
312
+ const state: Record<string, any> = {};
313
+ r.renderCall("grep", { pattern: "foo" }, theme, makeContext("a", state) as any);
314
+ const comp = r.renderResult(
315
+ "grep",
316
+ { pattern: "foo" },
317
+ { content: [{ type: "text", text: "a.ts:1:foo\nb.ts:2:foo\nc.ts:3:foo" }], details: { totalMatched: 3 } },
318
+ { isPartial: false, expanded: true },
319
+ theme,
320
+ { ...makeContext("a", state), isError: false } as any,
321
+ );
322
+ expect((comp as any).text).toContain("a.ts:1:foo");
323
+ expect((comp as any).text).toContain("b.ts:2:foo");
324
+ expect((comp as any).text).toContain("c.ts:3:foo");
325
+ });
326
+
327
+ test("expanded grouped call shows full output for owner", () => {
328
+ const r = new CompactRenderer();
329
+ const theme = makeTheme() as any;
330
+ const stateA: Record<string, any> = {};
331
+ const stateB: Record<string, any> = {};
332
+ r.renderCall("read", { file_path: "a.ts" }, theme, makeContext("a", stateA) as any);
333
+ r.renderCall("read", { file_path: "b.ts" }, theme, makeContext("b", stateB) as any);
334
+ const comp = r.renderResult(
335
+ "read",
336
+ { file_path: "b.ts" },
337
+ { content: [{ type: "text", text: "contentB_line1\ncontentB_line2" }] },
338
+ { isPartial: false, expanded: true },
339
+ theme,
340
+ { ...makeContext("b", stateB), isError: false } as any,
341
+ );
342
+ expect((comp as any).text).toContain("contentB_line1");
343
+ expect((comp as any).text).toContain("contentB_line2");
344
+ });
345
+
346
+ test("expanded partial result does not show output", () => {
347
+ const r = new CompactRenderer();
348
+ const theme = makeTheme() as any;
349
+ const state: Record<string, any> = {};
350
+ r.renderCall("bash", { command: "sleep 1" }, theme, makeContext("a", state) as any);
351
+ const comp = r.renderResult(
352
+ "bash",
353
+ { command: "sleep 1" },
354
+ { content: [{ type: "text", text: "intermediate" }] },
355
+ { isPartial: true, expanded: true },
356
+ theme,
357
+ { ...makeContext("a", state), isError: false } as any,
358
+ );
359
+ expect((comp as any).text).toBe("");
360
+ });
361
+
362
+ test("expanded error still shows error text", () => {
363
+ const r = new CompactRenderer();
364
+ const theme = makeTheme() as any;
365
+ r.renderCall("read", { file_path: "bad.ts" }, theme, makeContext("a") as any);
366
+ const comp = r.renderResult(
367
+ "read",
368
+ { file_path: "bad.ts" },
369
+ { content: [{ type: "text", text: "Error: not found" }] },
370
+ { isPartial: false, expanded: true },
371
+ theme,
372
+ { ...makeContext("a"), isError: true } as any,
373
+ );
374
+ expect((comp as any).text).toContain("Error: not found");
375
+ });
376
+ });
@@ -0,0 +1,122 @@
1
+ /**
2
+ * Ember TPS Meter — minimal tokens-per-second tracker
3
+ *
4
+ * Tracks output token rate during streaming and exposes the live value
5
+ * via getLiveTps() for the custom footer to render.
6
+ */
7
+
8
+ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
9
+
10
+ const STREAM_INTERVAL_MS = 500;
11
+
12
+ let streamStartMs = 0;
13
+ let firstTokenMs = 0;
14
+ let streamChars = 0;
15
+ let streamTokens = 0;
16
+ let tickTimer: ReturnType<typeof setInterval> | null = null;
17
+ let streaming = false;
18
+ let liveTps = 0;
19
+ let renderTrigger: (() => void) | undefined;
20
+
21
+ function now(): number {
22
+ return performance.now();
23
+ }
24
+
25
+ function tokEst(ch: number): number {
26
+ return (ch >>> 2) + ((ch & 3) > 0 ? 1 : 0);
27
+ }
28
+
29
+ function computeTps(): number {
30
+ const ref = firstTokenMs > 0 ? firstTokenMs : streamStartMs;
31
+ const elapsed = (now() - ref) / 1000;
32
+ return elapsed > 0.3 ? streamTokens / elapsed : 0;
33
+ }
34
+
35
+ function startTick(): void {
36
+ if (tickTimer) return;
37
+ tickTimer = setInterval(() => {
38
+ if (!streaming) {
39
+ stopTick();
40
+ return;
41
+ }
42
+ liveTps = computeTps();
43
+ renderTrigger?.();
44
+ }, STREAM_INTERVAL_MS);
45
+ }
46
+
47
+ function stopTick(): void {
48
+ if (tickTimer) {
49
+ clearInterval(tickTimer);
50
+ tickTimer = null;
51
+ }
52
+ }
53
+
54
+ export function getLiveTps(): number {
55
+ return liveTps;
56
+ }
57
+
58
+ export default function piEmberTps(pi: ExtensionAPI): void {
59
+ pi.on("message_start", async (event, ctx) => {
60
+ if (event.message.role !== "assistant") return;
61
+ streamStartMs = now();
62
+ firstTokenMs = 0;
63
+ streamChars = 0;
64
+ streamTokens = 0;
65
+ liveTps = 0;
66
+ streaming = true;
67
+ if (ctx.mode === "tui") {
68
+ renderTrigger = () => ctx.ui.setStatus("tps", undefined);
69
+ }
70
+ startTick();
71
+ });
72
+
73
+ pi.on("message_update", async (event) => {
74
+ if (event.message.role !== "assistant") return;
75
+ if (!event.assistantMessageEvent) return;
76
+ const evt = event.assistantMessageEvent;
77
+ if (evt.type === "text_delta" || evt.type === "thinking_delta") {
78
+ const d = evt.delta as string;
79
+ if (!d) return;
80
+ if (firstTokenMs === 0) firstTokenMs = now();
81
+ streamChars += d.length;
82
+ streamTokens = tokEst(streamChars);
83
+ }
84
+ });
85
+
86
+ pi.on("message_end", async (event) => {
87
+ if (event.message.role !== "assistant") return;
88
+ streaming = false;
89
+ stopTick();
90
+
91
+ const realOut = event.message?.usage?.output;
92
+ const tokens =
93
+ typeof realOut === "number" && realOut > 0 ? realOut : streamTokens;
94
+
95
+ const ref = firstTokenMs > 0 ? firstTokenMs : streamStartMs;
96
+ const elapsed = (now() - ref) / 1000;
97
+ if (elapsed < 0.1 || tokens === 0) {
98
+ liveTps = 0;
99
+ return;
100
+ }
101
+
102
+ liveTps = tokens / elapsed;
103
+ renderTrigger?.();
104
+ });
105
+
106
+ pi.on("agent_end", async () => {
107
+ streaming = false;
108
+ stopTick();
109
+ });
110
+
111
+ pi.on("session_start", async (_event, ctx) => {
112
+ streaming = false;
113
+ stopTick();
114
+ streamStartMs = 0;
115
+ firstTokenMs = 0;
116
+ streamChars = 0;
117
+ streamTokens = 0;
118
+ liveTps = 0;
119
+ renderTrigger = undefined;
120
+ ctx.ui.setStatus("tps", undefined);
121
+ });
122
+ }