@shanepadgett/tau-agent 0.28.1 → 0.30.0

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.
Files changed (66) hide show
  1. package/docs/context.md +14 -9
  2. package/docs/subagents.md +7 -5
  3. package/extensions/cache-diagnostics/index.ts +4 -5
  4. package/extensions/context/README.md +8 -5
  5. package/extensions/context/definitions.ts +28 -17
  6. package/extensions/context/evidence.ts +4 -3
  7. package/extensions/context/index.ts +57 -35
  8. package/extensions/context/panel.ts +10 -9
  9. package/extensions/context/projection.ts +141 -0
  10. package/extensions/context/state.ts +30 -0
  11. package/extensions/explore/README.md +1 -1
  12. package/extensions/explore/ast/read/hook.ts +5 -35
  13. package/extensions/explore/ast/read/policy.ts +0 -32
  14. package/extensions/explore/index.ts +2 -0
  15. package/extensions/explore/outline-injection.ts +151 -0
  16. package/extensions/explore/settings.ts +0 -8
  17. package/extensions/ideas/browser.ts +27 -16
  18. package/extensions/review/README.md +11 -0
  19. package/extensions/review/index.ts +135 -0
  20. package/extensions/review/model.ts +144 -0
  21. package/extensions/review/panel.ts +128 -0
  22. package/extensions/review/session.ts +106 -0
  23. package/extensions/runtime-context/README.md +1 -1
  24. package/extensions/runtime-context/index.ts +3 -66
  25. package/extensions/script-runner/README.md +7 -0
  26. package/extensions/script-runner/index.ts +275 -0
  27. package/extensions/stash/browser.ts +30 -18
  28. package/extensions/subagent/README.md +17 -4
  29. package/extensions/subagent/agents/context-sync.md +2 -2
  30. package/extensions/subagent/agents/scout.md +49 -52
  31. package/extensions/subagent/agents.ts +0 -1
  32. package/extensions/subagent/cmux-dashboard.ts +7 -3
  33. package/extensions/subagent/index.ts +105 -4
  34. package/extensions/subagent/panel.ts +124 -0
  35. package/extensions/subagent/render.ts +2 -1
  36. package/extensions/subagent/run.ts +9 -9
  37. package/extensions/subagent/runtime.ts +5 -5
  38. package/extensions/subagent/session-resource.ts +19 -127
  39. package/extensions/subagent/settings.ts +18 -0
  40. package/extensions/tau-help/help.md +12 -8
  41. package/extensions/working-memory/README.md +9 -0
  42. package/extensions/working-memory/checkpoint.ts +175 -0
  43. package/extensions/working-memory/index.ts +338 -0
  44. package/extensions/working-memory/memory.ts +266 -0
  45. package/extensions/working-memory/render.ts +178 -0
  46. package/extensions/working-memory/settings.ts +38 -0
  47. package/extensions/working-memory/state.ts +152 -0
  48. package/package.json +2 -2
  49. package/schemas/tau.schema.json +35 -55
  50. package/shared/context-messages.ts +19 -0
  51. package/shared/events.ts +5 -0
  52. package/shared/full-file-knowledge.ts +0 -1
  53. package/shared/injected-context.ts +2 -2
  54. package/shared/isolated-session.ts +151 -0
  55. package/shared/outline-injection.ts +56 -0
  56. package/extensions/context-pruning/README.md +0 -39
  57. package/extensions/context-pruning/index.ts +0 -382
  58. package/extensions/context-pruning/projection.ts +0 -60
  59. package/extensions/context-pruning/prune.ts +0 -199
  60. package/extensions/context-pruning/render.ts +0 -251
  61. package/extensions/context-pruning/settings.ts +0 -39
  62. package/extensions/subagent/agents/review.md +0 -68
  63. package/extensions/turn-budget/README.md +0 -14
  64. package/extensions/turn-budget/index.ts +0 -116
  65. package/extensions/turn-budget/settings.ts +0 -35
  66. package/shared/context-pruning-state.ts +0 -152
@@ -0,0 +1,151 @@
1
+ import {
2
+ createAgentSession,
3
+ DefaultResourceLoader,
4
+ getAgentDir,
5
+ ModelRuntime,
6
+ SessionManager,
7
+ type AgentSession,
8
+ type ExtensionContext,
9
+ type ToolDefinition,
10
+ } from "@earendil-works/pi-coding-agent";
11
+
12
+ type IsolatedSessionThinkingLevel = NonNullable<ExtensionContext["thinkingLevel"]>;
13
+
14
+ type SelectedModel = NonNullable<ExtensionContext["model"]>;
15
+ type SelectedProvider = NonNullable<ReturnType<ExtensionContext["modelRegistry"]["getProvider"]>>;
16
+
17
+ export interface IsolatedSessionInputs {
18
+ label: string;
19
+ extensionPaths: readonly string[];
20
+ cwd: string;
21
+ model: SelectedModel;
22
+ provider: SelectedProvider;
23
+ runtimeApiKey: string | undefined;
24
+ thinkingLevel: IsolatedSessionThinkingLevel;
25
+ tools: readonly string[];
26
+ customTools: readonly ToolDefinition[];
27
+ bindTarget: { mode: "print" } | { mode: "tui"; uiContext: ExtensionContext["ui"] };
28
+ }
29
+
30
+ export interface IsolatedSessionResource {
31
+ readonly session: AgentSession;
32
+ dispose(): Promise<void>;
33
+ }
34
+
35
+ export async function resolveIsolatedSessionModel(options: {
36
+ label: string;
37
+ preferredModel: string | undefined;
38
+ preferredThinkingLevel: IsolatedSessionThinkingLevel | undefined;
39
+ usePreferredThinkingAfterModelFallback: boolean;
40
+ ctx: ExtensionContext;
41
+ parentThinkingLevel: IsolatedSessionThinkingLevel;
42
+ signal: AbortSignal;
43
+ onWarning?: (warning: string) => void;
44
+ }): Promise<Pick<IsolatedSessionInputs, "model" | "provider" | "runtimeApiKey" | "thinkingLevel">> {
45
+ const { ctx, signal, onWarning } = options;
46
+ let model = ctx.model;
47
+ let thinkingLevel = options.parentThinkingLevel;
48
+ let preferredModelSelected = false;
49
+ let selectedAuth: Awaited<ReturnType<ExtensionContext["modelRegistry"]["getApiKeyAndHeaders"]>> | undefined;
50
+ if (options.preferredModel) {
51
+ const separator = options.preferredModel.indexOf("/");
52
+ const configured = ctx.modelRegistry.find(
53
+ options.preferredModel.slice(0, separator),
54
+ options.preferredModel.slice(separator + 1),
55
+ );
56
+ if (!configured) onWarning?.(`model ${options.preferredModel} is unavailable; using parent model`);
57
+ else {
58
+ const auth = await ctx.modelRegistry.getApiKeyAndHeaders(configured);
59
+ if (!auth.ok) onWarning?.(`model ${options.preferredModel} is unavailable: ${auth.error}; using parent model`);
60
+ else {
61
+ model = configured;
62
+ selectedAuth = auth;
63
+ preferredModelSelected = true;
64
+ }
65
+ }
66
+ }
67
+ if (options.preferredThinkingLevel && (preferredModelSelected || options.usePreferredThinkingAfterModelFallback)) {
68
+ const preferred = options.preferredThinkingLevel;
69
+ const mapped = model?.thinkingLevelMap?.[preferred];
70
+ const unsupported =
71
+ !model?.reasoning ||
72
+ mapped === null ||
73
+ ((preferred === "xhigh" || preferred === "max") && mapped === undefined);
74
+ if (unsupported)
75
+ onWarning?.(`thinking ${preferred} is unavailable for the selected model; using parent thinking`);
76
+ else thinkingLevel = preferred;
77
+ }
78
+ if (!model) throw new Error(`${options.label} startup failed: parent has no model`);
79
+ const auth = selectedAuth ?? (await ctx.modelRegistry.getApiKeyAndHeaders(model));
80
+ if (!auth.ok) throw new Error(`${options.label} startup failed: ${auth.error}`);
81
+ const provider = ctx.modelRegistry.getProvider(model.provider);
82
+ if (!provider) throw new Error(`${options.label} startup failed: provider ${model.provider} is unavailable`);
83
+ if (signal.aborted) throw new Error(`${options.label} startup aborted`);
84
+ return {
85
+ model,
86
+ provider,
87
+ runtimeApiKey:
88
+ auth.apiKey && provider.auth.apiKey && !ctx.modelRegistry.isUsingOAuth(model) ? auth.apiKey : undefined,
89
+ thinkingLevel,
90
+ };
91
+ }
92
+
93
+ export async function createIsolatedSessionResource(
94
+ inputs: IsolatedSessionInputs,
95
+ signal: AbortSignal,
96
+ ): Promise<IsolatedSessionResource> {
97
+ let session: AgentSession | undefined;
98
+ try {
99
+ if (signal.aborted) throw new Error(`${inputs.label} startup aborted`);
100
+ const modelRuntime = await ModelRuntime.create();
101
+ modelRuntime.registerNativeProvider(inputs.provider);
102
+ if (inputs.runtimeApiKey !== undefined)
103
+ await modelRuntime.setRuntimeApiKey(inputs.model.provider, inputs.runtimeApiKey, { allowNetwork: false });
104
+ if (signal.aborted) throw new Error(`${inputs.label} startup aborted`);
105
+ const resourceLoader = new DefaultResourceLoader({
106
+ cwd: inputs.cwd,
107
+ agentDir: getAgentDir(),
108
+ noExtensions: true,
109
+ additionalExtensionPaths: [...inputs.extensionPaths],
110
+ });
111
+ await resourceLoader.reload();
112
+ if (signal.aborted) throw new Error(`${inputs.label} startup aborted`);
113
+ const created = await createAgentSession({
114
+ cwd: inputs.cwd,
115
+ model: inputs.model,
116
+ modelRuntime,
117
+ thinkingLevel: inputs.thinkingLevel,
118
+ tools: [...inputs.tools],
119
+ excludeTools: ["subagent"],
120
+ resourceLoader,
121
+ customTools: [...inputs.customTools],
122
+ sessionManager: SessionManager.inMemory(inputs.cwd),
123
+ });
124
+ session = created.session;
125
+ if (signal.aborted) throw new Error(`${inputs.label} startup aborted`);
126
+ await session.bindExtensions(inputs.bindTarget);
127
+ if (signal.aborted) throw new Error(`${inputs.label} startup aborted`);
128
+ const active = session.getActiveToolNames().sort();
129
+ const expected = [...inputs.tools].sort();
130
+ if (active.join("\0") !== expected.join("\0") || active.includes("subagent")) {
131
+ const missing = expected.filter((tool) => !active.includes(tool));
132
+ throw new Error(
133
+ `${inputs.label} startup failed: unavailable tools: ${missing.join(", ") || "active tool mismatch"}`,
134
+ );
135
+ }
136
+ let disposed = false;
137
+ return {
138
+ session,
139
+ async dispose() {
140
+ if (disposed) return;
141
+ disposed = true;
142
+ if (session?.isStreaming) await session.abort().catch(() => undefined);
143
+ session?.dispose();
144
+ },
145
+ };
146
+ } catch (error) {
147
+ if (session?.isStreaming) await session.abort().catch(() => undefined);
148
+ session?.dispose();
149
+ throw error;
150
+ }
151
+ }
@@ -0,0 +1,56 @@
1
+ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
2
+ import { emitTauEvent, onTauEvent } from "./events.js";
3
+
4
+ export interface OutlineInjectionDetails {
5
+ v: 1;
6
+ rowId: string;
7
+ path: string;
8
+ cwd: string;
9
+ batchId: string;
10
+ }
11
+
12
+ export interface PreparedOutlineInjection {
13
+ customType: "tau.explore.outline";
14
+ content: string;
15
+ display: true;
16
+ details: OutlineInjectionDetails;
17
+ }
18
+
19
+ export interface OutlineInjectionRequest {
20
+ cwd: string;
21
+ batchId: string;
22
+ paths: readonly string[];
23
+ signal: AbortSignal | undefined;
24
+ isLifecycleCurrent(): boolean;
25
+ }
26
+
27
+ export interface OutlineInjectionResponse {
28
+ messages: PreparedOutlineInjection[];
29
+ warnings: string[];
30
+ }
31
+
32
+ export type OutlineInjectionProvider = (request: OutlineInjectionRequest) => Promise<OutlineInjectionResponse>;
33
+
34
+ export function registerOutlineInjectionProvider(pi: ExtensionAPI, provider: OutlineInjectionProvider): () => void {
35
+ return onTauEvent(pi, "shared.outline-injection-provider", "tau:outline-injection.requested", (event) => {
36
+ event.accept(provider(event.request));
37
+ });
38
+ }
39
+
40
+ export async function requestOutlineInjections(
41
+ pi: Pick<ExtensionAPI, "events">,
42
+ request: OutlineInjectionRequest,
43
+ ): Promise<OutlineInjectionResponse> {
44
+ let response: Promise<OutlineInjectionResponse> | undefined;
45
+ emitTauEvent(pi, "tau:outline-injection.requested", {
46
+ request,
47
+ accept(candidate) {
48
+ response ??= candidate;
49
+ },
50
+ });
51
+ if (response) return response;
52
+ return {
53
+ messages: [],
54
+ warnings: request.paths.map((path) => `${path}: Explore outline provider is unavailable`),
55
+ };
56
+ }
@@ -1,39 +0,0 @@
1
- # Context Pruning
2
-
3
- Context Pruning gives the agent direct control over its future model context. It does not delete or rewrite the saved conversation.
4
-
5
- `context_prune` creates a hard checkpoint. Everything before that checkpoint leaves future model input unless the agent explicitly carries it forward. Messages after the checkpoint remain unchanged.
6
-
7
- Before calling the tool, the agent writes visible prose containing the durable conclusions, user constraints, conditional relevance, and next action that must survive. That prose is part of the checkpoint turn.
8
-
9
- ## Selections
10
-
11
- The tool accepts three required lists:
12
-
13
- - `keepToolCalls` retains exact earlier tool exchanges by tool-call ID. Parallel calls remain independently selectable: retaining one call does not retain its siblings.
14
- - `keepFiles` reads each selected file from disk and carries its complete current contents forward as a fresh autoread snapshot. Each successful snapshot appears as its own autoread marker below the compact checkpoint result. It does not require an earlier complete read.
15
- - `deferFiles` carries forward a short advisory note explaining why a file is irrelevant now and when to reconsider it.
16
-
17
- Duplicate selections are collapsed. A file selected in both `keepFiles` and `deferFiles` is kept. Missing, unreadable, or otherwise unsnapshotable files produce warnings in the successful tool result; they do not block the checkpoint or other file snapshots.
18
-
19
- Context Pruning does not impose a minimum token saving and does not reject a checkpoint because old context has unusual bookkeeping. Pi supplies the current provider context. Tau removes unretained tool calls and their matching results with the same ID, drops other pre-checkpoint messages, and leaves the checkpoint turn and later messages intact.
20
-
21
- ## Automatic and manual requests
22
-
23
- Tau checks active context size after tool-using turns and can send progressively stronger private instructions from `nudgeInstructions`. Hints occur at fixed token intervals rather than percentages of the model's context window. After a checkpoint, Tau suppresses boundaries already crossed by the resulting context and resumes at the next boundary. Branch navigation and compaction reconstruct that state from the active branch.
24
-
25
- Run `/prune` with no arguments to ask the agent to create a checkpoint and continue its task immediately.
26
-
27
- ## Branches, compaction, and display
28
-
29
- Checkpoints belong to the active branch. Switching branches rebuilds the latest checkpoint, deferred-file state, automatic-hint baseline, and warning-colored pruned rows from that branch.
30
-
31
- Normal Pi compaction remains independent. When a compaction no longer includes an old checkpoint turn, that checkpoint has nothing left to filter.
32
-
33
- ## Settings
34
-
35
- Settings live under `extensions.contextPruning` in Tau settings.
36
-
37
- - `enabled`: enables the tool, `/prune`, projection, markers, and branch replay. Defaults to `true`.
38
- - `nudgeEveryTokens`: active-context interval between automatic hints. Defaults to `30000`, producing the default instruction ladder at 30k, 60k, and 90k tokens.
39
- - `nudgeInstructions`: ordered list of one through five nonempty instructions. Later reminders repeat the final instruction. Defaults to three escalating instructions.
@@ -1,382 +0,0 @@
1
- import {
2
- defineTool,
3
- type ExtensionAPI,
4
- type ExtensionContext,
5
- type SessionEntry,
6
- } from "@earendil-works/pi-coding-agent";
7
- import {
8
- replayContextPruningState,
9
- setContextPruningEnabled,
10
- type ContextPruneDetailsV2,
11
- } from "../../shared/context-pruning-state.ts";
12
- import { emitTauEvent, onTauEvent } from "../../shared/events.ts";
13
- import { loadTauExtensionSettings } from "../../shared/settings/load.ts";
14
- import { createToolRowStateStore } from "../../shared/tool-row-state.ts";
15
- import { contextPruneParameters, executeContextPrune } from "./prune.ts";
16
- import { projectContext } from "./projection.ts";
17
- import {
18
- parseContextPruningNudgeDetailsV3,
19
- renderContextPruneCall,
20
- renderContextPruneResult,
21
- renderContextPruningNudge,
22
- type ContextPruningNudgeDetailsV3,
23
- } from "./render.ts";
24
- import contextPruningSettings from "./settings.ts";
25
-
26
- const TOOL_DESCRIPTION =
27
- "Create a hard context checkpoint after broad exploration converges, when a context-pruning nudge directs it, or when stale evidence has accumulated. Everything before the checkpoint is removed from future model context unless selected for retention. Immediately before calling, state durable conclusions, conditional relevance, and the next action in visible prose.";
28
- const NUDGE_MESSAGE_TYPE = "tau.context-pruning.nudge";
29
- const NUDGE_BASELINE_ENTRY_TYPE = "tau.context-pruning.nudge-baseline";
30
-
31
- interface NudgeState {
32
- anchorToolCallId: string | undefined;
33
- suppressedThroughTokens: number | undefined;
34
- highestBoundaryTokens: number;
35
- highestTier: number;
36
- terminalTierReached: boolean;
37
- }
38
-
39
- export default function contextPruningExtension(pi: ExtensionAPI): void {
40
- let enabled = false;
41
- let lifecycleGeneration = 0;
42
- let nudgeEveryTokens = contextPruningSettings.defaults.nudgeEveryTokens;
43
- let nudgeInstructions = contextPruningSettings.defaults.nudgeInstructions;
44
- let toolRegistered = false;
45
- let commandRegistered = false;
46
- let visualRows = new Set<string>();
47
- let nudgeState: NudgeState = {
48
- anchorToolCallId: undefined,
49
- suppressedThroughTokens: 0,
50
- highestBoundaryTokens: 0,
51
- highestTier: 0,
52
- terminalTierReached: false,
53
- };
54
- const rowState = createToolRowStateStore(pi, "context-pruning.tool-row-state");
55
- pi.registerMessageRenderer<ContextPruningNudgeDetailsV3>(NUDGE_MESSAGE_TYPE, (message, _options, theme) =>
56
- renderContextPruningNudge(message.details, theme),
57
- );
58
-
59
- const pushVisualSnapshot = () => {
60
- emitTauEvent(pi, "tau:tool-row-state.snapshot", {
61
- states: [...visualRows].map((rowId) => ({ rowId, state: "pruned" as const })),
62
- });
63
- };
64
- onTauEvent(
65
- pi,
66
- "context-pruning.tool-row-state-producer",
67
- "tau:tool-row-state.snapshot.requested",
68
- pushVisualSnapshot,
69
- );
70
-
71
- const clearEphemeralState = () => {
72
- lifecycleGeneration += 1;
73
- };
74
- const setContextPruneToolActive = (active: boolean) => {
75
- if (!toolRegistered) return;
76
- const activeTools = pi.getActiveTools();
77
- const currentlyActive = activeTools.includes("context_prune");
78
- if (active === currentlyActive) return;
79
- pi.setActiveTools(
80
- active ? [...activeTools, "context_prune"] : activeTools.filter((toolName) => toolName !== "context_prune"),
81
- );
82
- };
83
- const syncBranchState = (ctx: ExtensionContext) => {
84
- const state = replayContextPruningState(ctx.sessionManager.getBranch(), enabled);
85
- visualRows = new Set([...state.prunedToolCallIds, ...state.prunedAutoreadRowIds]);
86
- nudgeState = reconstructNudgeState(ctx.sessionManager.getBranch(), state.latestAnchorToolCallId);
87
- pushVisualSnapshot();
88
- };
89
-
90
- pi.on("session_start", async (_event, ctx) => {
91
- clearEphemeralState();
92
- enabled = false;
93
- setContextPruneToolActive(false);
94
- visualRows.clear();
95
- setContextPruningEnabled(false);
96
- pushVisualSnapshot();
97
- const generation = lifecycleGeneration;
98
- const settings = await loadTauExtensionSettings(ctx, contextPruningSettings);
99
- if (generation !== lifecycleGeneration) return;
100
- enabled = settings.enabled;
101
- nudgeEveryTokens = settings.nudgeEveryTokens;
102
- nudgeInstructions = settings.nudgeInstructions;
103
- setContextPruningEnabled(enabled);
104
- if (enabled && !toolRegistered) {
105
- pi.registerTool(
106
- defineTool<typeof contextPruneParameters, ContextPruneDetailsV2>({
107
- name: "context_prune",
108
- label: "context_prune",
109
- description: TOOL_DESCRIPTION,
110
- promptSnippet:
111
- "Prune substantial stale tool evidence after stating durable conclusions and the next action",
112
- promptGuidelines: [
113
- "Use context_prune after broad exploration converges, when a context-pruning nudge directs it, or when substantial irrelevant evidence has accumulated.",
114
- "A final-tier context-pruning nudge means preserve durable conclusions and prune before further tool work.",
115
- "Everything before context_prune leaves future model context unless selected in keepFiles or keepToolCalls, so preserve durable conclusions, user constraints, conditional relevance, and the next action in visible prose immediately before calling it.",
116
- ],
117
- parameters: contextPruneParameters,
118
- executionMode: "sequential",
119
- async execute(toolCallId, params, signal, _onUpdate, executionContext) {
120
- const execution = await executeContextPrune({
121
- toolCallId,
122
- params,
123
- signal,
124
- ctx: executionContext,
125
- generation: lifecycleGeneration,
126
- currentGeneration: () => lifecycleGeneration,
127
- });
128
- for (const autoread of execution.autoreads) {
129
- pi.sendMessage(autoread, { deliverAs: "steer" });
130
- }
131
- return execution.result;
132
- },
133
- renderCall(args, theme, context) {
134
- return renderContextPruneCall(args, theme, {
135
- rowState,
136
- rowId: context.toolCallId,
137
- invalidate: context.invalidate,
138
- lastComponent: context.lastComponent,
139
- });
140
- },
141
- renderResult(result, options, theme, context) {
142
- return renderContextPruneResult(result, options.expanded, theme, context.lastComponent);
143
- },
144
- }),
145
- );
146
- toolRegistered = true;
147
- }
148
- setContextPruneToolActive(enabled);
149
- if (enabled && !commandRegistered) {
150
- pi.registerCommand("prune", {
151
- description: "Ask the agent to create a context-pruning anchor and continue its task",
152
- async handler(args, commandContext) {
153
- if (!enabled) {
154
- commandContext.ui.notify("Context pruning is disabled.", "info");
155
- return;
156
- }
157
- if (args.trim().length > 0) {
158
- commandContext.ui.notify("Usage: /prune", "info");
159
- return;
160
- }
161
- const anchorToolCallId = replayContextPruningState(
162
- commandContext.sessionManager.getBranch(),
163
- true,
164
- ).latestAnchorToolCallId;
165
- pi.sendMessage<ContextPruningNudgeDetailsV3>(
166
- {
167
- customType: NUDGE_MESSAGE_TYPE,
168
- content: manualPruneSteeringMessage(),
169
- display: true,
170
- details: {
171
- v: 3,
172
- kind: "manual",
173
- tokens: null,
174
- boundaryTokens: null,
175
- reminder: null,
176
- tier: null,
177
- tierCount: null,
178
- tierFloor: null,
179
- anchorToolCallId: anchorToolCallId ?? null,
180
- suppressedThroughTokens: null,
181
- },
182
- },
183
- { deliverAs: "steer", triggerTurn: true },
184
- );
185
- },
186
- });
187
- commandRegistered = true;
188
- }
189
- syncBranchState(ctx);
190
- });
191
-
192
- pi.on("session_tree", (_event, ctx) => {
193
- clearEphemeralState();
194
- syncBranchState(ctx);
195
- });
196
- pi.on("session_compact", (_event, ctx) => {
197
- clearEphemeralState();
198
- syncBranchState(ctx);
199
- });
200
- pi.on("session_shutdown", () => {
201
- clearEphemeralState();
202
- enabled = false;
203
- setContextPruneToolActive(false);
204
- visualRows.clear();
205
- nudgeState = {
206
- anchorToolCallId: undefined,
207
- suppressedThroughTokens: 0,
208
- highestBoundaryTokens: 0,
209
- highestTier: 0,
210
- terminalTierReached: false,
211
- };
212
- setContextPruningEnabled(false);
213
- pushVisualSnapshot();
214
- });
215
-
216
- pi.on("turn_end", (event, ctx) => {
217
- if (!enabled || event.toolResults.length === 0) return undefined;
218
- const usage = ctx.getContextUsage();
219
- if (!usage || usage.tokens === null || !Number.isFinite(usage.tokens)) return undefined;
220
- const tokens = Math.max(0, Math.floor(usage.tokens));
221
- const activeAnchor = replayContextPruningState(ctx.sessionManager.getBranch(), true).latestAnchorToolCallId;
222
- if (activeAnchor !== nudgeState.anchorToolCallId) {
223
- nudgeState = reconstructNudgeState(ctx.sessionManager.getBranch(), activeAnchor);
224
- }
225
- if (activeAnchor !== undefined && nudgeState.suppressedThroughTokens === undefined) {
226
- const suppressedThroughTokens = Math.floor(tokens / nudgeEveryTokens) * nudgeEveryTokens;
227
- pi.appendEntry(NUDGE_BASELINE_ENTRY_TYPE, {
228
- v: 2,
229
- anchorToolCallId: activeAnchor,
230
- suppressedThroughTokens,
231
- });
232
- nudgeState.suppressedThroughTokens = suppressedThroughTokens;
233
- nudgeState.highestBoundaryTokens = suppressedThroughTokens;
234
- return undefined;
235
- }
236
- const reminder = Math.floor(tokens / nudgeEveryTokens);
237
- if (reminder < 1) return undefined;
238
- const boundaryTokens = reminder * nudgeEveryTokens;
239
- if (boundaryTokens <= nudgeState.highestBoundaryTokens) return undefined;
240
- const tierCount = nudgeInstructions.length;
241
- const tierFloor = nudgeState.terminalTierReached ? tierCount : Math.min(nudgeState.highestTier, tierCount);
242
- const tier = Math.max(Math.min(reminder, tierCount), tierFloor);
243
- const instruction = nudgeInstructions[tier - 1] ?? nudgeInstructions[0];
244
- const details: ContextPruningNudgeDetailsV3 = {
245
- v: 3,
246
- kind: "automatic",
247
- tokens,
248
- boundaryTokens,
249
- reminder,
250
- tier,
251
- tierCount,
252
- tierFloor,
253
- anchorToolCallId: activeAnchor ?? null,
254
- suppressedThroughTokens: nudgeState.suppressedThroughTokens ?? 0,
255
- };
256
- pi.sendMessage<ContextPruningNudgeDetailsV3>(
257
- {
258
- customType: NUDGE_MESSAGE_TYPE,
259
- content: automaticPruneSteeringMessage(instruction, tier === tierCount),
260
- display: true,
261
- details,
262
- },
263
- { deliverAs: "steer" },
264
- );
265
- nudgeState.highestBoundaryTokens = boundaryTokens;
266
- nudgeState.highestTier = Math.max(nudgeState.highestTier, tier);
267
- nudgeState.terminalTierReached ||= tier === tierCount;
268
- return undefined;
269
- });
270
-
271
- pi.on("context", (event, ctx) => {
272
- if (!enabled) return undefined;
273
- const state = replayContextPruningState(ctx.sessionManager.getBranch(), true);
274
- const messages = projectContext(event.messages, state);
275
- const nextRows = new Set([...state.prunedToolCallIds, ...state.prunedAutoreadRowIds]);
276
- if (!setsEqual(visualRows, nextRows)) {
277
- visualRows = nextRows;
278
- pushVisualSnapshot();
279
- }
280
- return { messages };
281
- });
282
- }
283
-
284
- function setsEqual(left: ReadonlySet<string>, right: ReadonlySet<string>): boolean {
285
- if (left.size !== right.size) return false;
286
- for (const item of left) if (!right.has(item)) return false;
287
- return true;
288
- }
289
-
290
- function reconstructNudgeState(branch: readonly SessionEntry[], anchorToolCallId: string | undefined): NudgeState {
291
- let suppressedThroughTokens = anchorToolCallId === undefined ? 0 : undefined;
292
- let highestBoundaryTokens = 0;
293
- let highestTier = 0;
294
- let terminalTierReached = false;
295
- let anchorResultIndex = -1;
296
- if (anchorToolCallId !== undefined) {
297
- anchorResultIndex = branch.findIndex(
298
- (entry) =>
299
- entry.type === "message" &&
300
- entry.message.role === "toolResult" &&
301
- entry.message.toolName === "context_prune" &&
302
- entry.message.toolCallId === anchorToolCallId,
303
- );
304
- }
305
- for (let index = 0; index < branch.length; index += 1) {
306
- const entry = branch[index];
307
- if (!entry) continue;
308
- if (entry.type === "custom" && entry.customType === NUDGE_BASELINE_ENTRY_TYPE) {
309
- const baseline = parseNudgeBaseline(entry.data);
310
- if (
311
- baseline &&
312
- suppressedThroughTokens === undefined &&
313
- index > anchorResultIndex &&
314
- baseline.anchorToolCallId === anchorToolCallId
315
- ) {
316
- suppressedThroughTokens = baseline.suppressedThroughTokens;
317
- highestBoundaryTokens = baseline.suppressedThroughTokens;
318
- }
319
- continue;
320
- }
321
- if (entry.type !== "custom_message" || entry.customType !== NUDGE_MESSAGE_TYPE) continue;
322
- const details = parseContextPruningNudgeDetailsV3(entry.details);
323
- if (
324
- !details ||
325
- details.kind !== "automatic" ||
326
- index <= anchorResultIndex ||
327
- details.anchorToolCallId !== (anchorToolCallId ?? null) ||
328
- details.boundaryTokens === null ||
329
- details.suppressedThroughTokens === null ||
330
- (suppressedThroughTokens !== undefined && details.suppressedThroughTokens !== suppressedThroughTokens)
331
- )
332
- continue;
333
- const expectedTierFloor: number = terminalTierReached
334
- ? details.tierCount
335
- : Math.min(highestTier, details.tierCount);
336
- if (details.boundaryTokens <= highestBoundaryTokens || details.tierFloor !== expectedTierFloor) continue;
337
- highestBoundaryTokens = Math.max(highestBoundaryTokens, details.boundaryTokens);
338
- highestTier = Math.max(highestTier, details.tier);
339
- terminalTierReached ||= details.tier === details.tierCount;
340
- suppressedThroughTokens = details.suppressedThroughTokens;
341
- }
342
- return { anchorToolCallId, suppressedThroughTokens, highestBoundaryTokens, highestTier, terminalTierReached };
343
- }
344
-
345
- function parseNudgeBaseline(
346
- value: unknown,
347
- ): { v: 2; anchorToolCallId: string; suppressedThroughTokens: number } | undefined {
348
- if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined;
349
- const record = value as Record<string, unknown>;
350
- if (
351
- Object.keys(record).length !== 3 ||
352
- !Object.hasOwn(record, "v") ||
353
- !Object.hasOwn(record, "anchorToolCallId") ||
354
- !Object.hasOwn(record, "suppressedThroughTokens") ||
355
- record.v !== 2 ||
356
- typeof record.anchorToolCallId !== "string" ||
357
- record.anchorToolCallId.length === 0 ||
358
- typeof record.suppressedThroughTokens !== "number" ||
359
- !Number.isSafeInteger(record.suppressedThroughTokens) ||
360
- record.suppressedThroughTokens < 0
361
- )
362
- return undefined;
363
- return {
364
- v: 2,
365
- anchorToolCallId: record.anchorToolCallId,
366
- suppressedThroughTokens: record.suppressedThroughTokens,
367
- };
368
- }
369
-
370
- function automaticPruneSteeringMessage(instruction: string, finalTier: boolean): string {
371
- const silent =
372
- "Internal context-management instruction. Follow it silently. Do not mention or acknowledge context-token counts, prune messages, or internal context management.";
373
- const protocol =
374
- "When pruning, first preserve durable conclusions, user constraints, conditional relevance, and the next action in visible prose, then call context_prune.";
375
- return finalTier
376
- ? `${silent} ${instruction} This is the final reminder tier. Create a context anchor before further tool work. ${protocol}`
377
- : `${silent} ${instruction} ${protocol}`;
378
- }
379
-
380
- function manualPruneSteeringMessage(): string {
381
- return "Internal context-management instruction. Follow it silently without mentioning this request. Create a hard context checkpoint with context_prune, then continue unfinished work. First preserve durable conclusions, user constraints, conditional relevance, and the next action in visible prose.";
382
- }
@@ -1,60 +0,0 @@
1
- import type { ContextEvent } from "@earendil-works/pi-coding-agent";
2
- import type { ActiveContextPruningState } from "../../shared/context-pruning-state.ts";
3
-
4
- type ContextMessage = ContextEvent["messages"][number];
5
-
6
- export function projectContext(
7
- messages: readonly ContextMessage[],
8
- state: ActiveContextPruningState,
9
- ): ContextMessage[] {
10
- if (state.latestAnchorToolCallId === undefined) return [...messages];
11
- let anchorIndex = -1;
12
- for (let index = messages.length - 1; index >= 0; index -= 1) {
13
- const message = messages[index];
14
- if (
15
- message?.role === "assistant" &&
16
- message.content.some(
17
- (block) =>
18
- block.type === "toolCall" && block.id === state.latestAnchorToolCallId && block.name === "context_prune",
19
- )
20
- ) {
21
- anchorIndex = index;
22
- break;
23
- }
24
- }
25
- if (anchorIndex < 0) return [...messages];
26
- const retainedCallNames = new Map<string, string>();
27
- const retainedResultNames = new Map<string, string>();
28
- for (let index = 0; index < anchorIndex; index += 1) {
29
- const message = messages[index];
30
- if (message?.role === "assistant") {
31
- for (const block of message.content) {
32
- if (block.type === "toolCall" && state.retainedToolCallIds.has(block.id)) {
33
- retainedCallNames.set(block.id, block.name);
34
- }
35
- }
36
- } else if (message?.role === "toolResult" && state.retainedToolCallIds.has(message.toolCallId)) {
37
- retainedResultNames.set(message.toolCallId, message.toolName);
38
- }
39
- }
40
- const retainableToolCallIds = new Set(
41
- [...retainedCallNames].flatMap(([id, name]) => (retainedResultNames.get(id) === name ? [id] : [])),
42
- );
43
-
44
- const projected: ContextMessage[] = [];
45
- for (let index = 0; index < anchorIndex; index += 1) {
46
- const message = messages[index];
47
- if (!message) continue;
48
- if (message.role === "toolResult") {
49
- if (retainableToolCallIds.has(message.toolCallId)) projected.push(message);
50
- continue;
51
- }
52
- if (message.role !== "assistant") continue;
53
- const content = message.content.filter(
54
- (block) => block.type === "toolCall" && retainableToolCallIds.has(block.id),
55
- );
56
- if (content.length > 0) projected.push({ ...message, content });
57
- }
58
- projected.push(...messages.slice(anchorIndex));
59
- return projected;
60
- }