omegon 0.6.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 (160) hide show
  1. package/.gitattributes +3 -0
  2. package/AGENTS.md +16 -0
  3. package/LICENSE +15 -0
  4. package/README.md +289 -0
  5. package/bin/pi.mjs +30 -0
  6. package/extensions/00-secrets/index.ts +1126 -0
  7. package/extensions/01-auth/auth.ts +401 -0
  8. package/extensions/01-auth/index.ts +289 -0
  9. package/extensions/auto-compact.ts +42 -0
  10. package/extensions/bootstrap/deps.ts +291 -0
  11. package/extensions/bootstrap/index.ts +811 -0
  12. package/extensions/chronos/chronos.sh +487 -0
  13. package/extensions/chronos/index.ts +148 -0
  14. package/extensions/cleave/assessment.ts +754 -0
  15. package/extensions/cleave/bridge.ts +31 -0
  16. package/extensions/cleave/conflicts.ts +250 -0
  17. package/extensions/cleave/dispatcher.ts +808 -0
  18. package/extensions/cleave/guardrails.ts +426 -0
  19. package/extensions/cleave/index.ts +3121 -0
  20. package/extensions/cleave/lifecycle-emitter.ts +20 -0
  21. package/extensions/cleave/openspec.ts +811 -0
  22. package/extensions/cleave/planner.ts +260 -0
  23. package/extensions/cleave/review.ts +579 -0
  24. package/extensions/cleave/skills.ts +355 -0
  25. package/extensions/cleave/types.ts +261 -0
  26. package/extensions/cleave/workspace.ts +861 -0
  27. package/extensions/cleave/worktree.ts +243 -0
  28. package/extensions/core-renderers.ts +253 -0
  29. package/extensions/dashboard/context-gauge.ts +58 -0
  30. package/extensions/dashboard/file-watch.ts +14 -0
  31. package/extensions/dashboard/footer.ts +1145 -0
  32. package/extensions/dashboard/git.ts +185 -0
  33. package/extensions/dashboard/index.ts +478 -0
  34. package/extensions/dashboard/memory-audit.ts +34 -0
  35. package/extensions/dashboard/overlay-data.ts +705 -0
  36. package/extensions/dashboard/overlay.ts +365 -0
  37. package/extensions/dashboard/render-utils.ts +54 -0
  38. package/extensions/dashboard/types.ts +191 -0
  39. package/extensions/dashboard/uri-helper.ts +45 -0
  40. package/extensions/debug.ts +69 -0
  41. package/extensions/defaults.ts +282 -0
  42. package/extensions/design-tree/dashboard-state.ts +161 -0
  43. package/extensions/design-tree/design-card.ts +362 -0
  44. package/extensions/design-tree/index.ts +2130 -0
  45. package/extensions/design-tree/lifecycle-emitter.ts +41 -0
  46. package/extensions/design-tree/tree.ts +1607 -0
  47. package/extensions/design-tree/types.ts +163 -0
  48. package/extensions/distill.ts +127 -0
  49. package/extensions/effort/index.ts +395 -0
  50. package/extensions/effort/tiers.ts +146 -0
  51. package/extensions/effort/types.ts +105 -0
  52. package/extensions/lib/git-state.ts +227 -0
  53. package/extensions/lib/local-models.ts +157 -0
  54. package/extensions/lib/model-preferences.ts +51 -0
  55. package/extensions/lib/model-routing.ts +720 -0
  56. package/extensions/lib/operator-fallback.ts +205 -0
  57. package/extensions/lib/operator-profile.ts +360 -0
  58. package/extensions/lib/slash-command-bridge.ts +253 -0
  59. package/extensions/lib/typebox-helpers.ts +16 -0
  60. package/extensions/local-inference/index.ts +727 -0
  61. package/extensions/mcp-bridge/README.md +220 -0
  62. package/extensions/mcp-bridge/index.ts +951 -0
  63. package/extensions/mcp-bridge/lib.ts +365 -0
  64. package/extensions/mcp-bridge/mcp.json +3 -0
  65. package/extensions/mcp-bridge/package.json +11 -0
  66. package/extensions/model-budget.ts +752 -0
  67. package/extensions/offline-driver.ts +403 -0
  68. package/extensions/openspec/archive-gate.ts +164 -0
  69. package/extensions/openspec/branch-cleanup.ts +64 -0
  70. package/extensions/openspec/dashboard-state.ts +50 -0
  71. package/extensions/openspec/index.ts +1917 -0
  72. package/extensions/openspec/lifecycle-emitter.ts +65 -0
  73. package/extensions/openspec/lifecycle-files.ts +70 -0
  74. package/extensions/openspec/lifecycle.ts +50 -0
  75. package/extensions/openspec/reconcile.ts +187 -0
  76. package/extensions/openspec/spec.ts +1385 -0
  77. package/extensions/openspec/types.ts +98 -0
  78. package/extensions/project-memory/DESIGN-global-mind.md +198 -0
  79. package/extensions/project-memory/README.md +202 -0
  80. package/extensions/project-memory/api-types.ts +382 -0
  81. package/extensions/project-memory/compaction-policy.ts +29 -0
  82. package/extensions/project-memory/core.ts +164 -0
  83. package/extensions/project-memory/embeddings.ts +230 -0
  84. package/extensions/project-memory/extraction-v2.ts +861 -0
  85. package/extensions/project-memory/factstore.ts +2177 -0
  86. package/extensions/project-memory/index.ts +3459 -0
  87. package/extensions/project-memory/injection-metrics.ts +91 -0
  88. package/extensions/project-memory/jsonl-io.ts +12 -0
  89. package/extensions/project-memory/lifecycle.ts +331 -0
  90. package/extensions/project-memory/migration.ts +293 -0
  91. package/extensions/project-memory/package.json +9 -0
  92. package/extensions/project-memory/sci-renderers.ts +7 -0
  93. package/extensions/project-memory/template.ts +103 -0
  94. package/extensions/project-memory/triggers.ts +52 -0
  95. package/extensions/project-memory/types.ts +102 -0
  96. package/extensions/render/composition/fonts/Inter-Bold.ttf +0 -0
  97. package/extensions/render/composition/fonts/Inter-Regular.ttf +0 -0
  98. package/extensions/render/composition/fonts/Tomorrow-Bold.ttf +0 -0
  99. package/extensions/render/composition/fonts/Tomorrow-Regular.ttf +0 -0
  100. package/extensions/render/composition/package-lock.json +534 -0
  101. package/extensions/render/composition/package.json +22 -0
  102. package/extensions/render/composition/render.mjs +246 -0
  103. package/extensions/render/composition/test-comp.tsx +87 -0
  104. package/extensions/render/composition/types.ts +24 -0
  105. package/extensions/render/excalidraw/UPSTREAM.md +81 -0
  106. package/extensions/render/excalidraw/elements.ts +764 -0
  107. package/extensions/render/excalidraw/index.ts +66 -0
  108. package/extensions/render/excalidraw/types.ts +223 -0
  109. package/extensions/render/excalidraw-renderer/pyproject.toml +8 -0
  110. package/extensions/render/excalidraw-renderer/render_excalidraw.py +182 -0
  111. package/extensions/render/excalidraw-renderer/render_template.html +59 -0
  112. package/extensions/render/index.ts +830 -0
  113. package/extensions/render/native-diagrams/index.ts +57 -0
  114. package/extensions/render/native-diagrams/motifs.ts +542 -0
  115. package/extensions/render/native-diagrams/raster.ts +8 -0
  116. package/extensions/render/native-diagrams/scene.ts +75 -0
  117. package/extensions/render/native-diagrams/spec.ts +204 -0
  118. package/extensions/render/native-diagrams/svg.ts +116 -0
  119. package/extensions/sci-ui.ts +304 -0
  120. package/extensions/session-log.ts +174 -0
  121. package/extensions/shared-state.ts +146 -0
  122. package/extensions/spinner-verbs.ts +91 -0
  123. package/extensions/style.ts +281 -0
  124. package/extensions/terminal-title.ts +191 -0
  125. package/extensions/tool-profile/index.ts +291 -0
  126. package/extensions/tool-profile/profiles.ts +290 -0
  127. package/extensions/types.d.ts +9 -0
  128. package/extensions/vault/index.ts +185 -0
  129. package/extensions/version-check.ts +90 -0
  130. package/extensions/view/index.ts +859 -0
  131. package/extensions/view/uri-resolver.ts +148 -0
  132. package/extensions/web-search/index.ts +182 -0
  133. package/extensions/web-search/providers.ts +121 -0
  134. package/extensions/web-ui/index.ts +110 -0
  135. package/extensions/web-ui/server.ts +265 -0
  136. package/extensions/web-ui/state.ts +462 -0
  137. package/extensions/web-ui/static/index.html +145 -0
  138. package/extensions/web-ui/types.ts +284 -0
  139. package/package.json +76 -0
  140. package/prompts/init.md +75 -0
  141. package/prompts/new-repo.md +54 -0
  142. package/prompts/oci-login.md +56 -0
  143. package/prompts/status.md +50 -0
  144. package/settings.json +4 -0
  145. package/skills/cleave/SKILL.md +218 -0
  146. package/skills/git/SKILL.md +209 -0
  147. package/skills/git/_reference/ci-validation.md +204 -0
  148. package/skills/oci/SKILL.md +338 -0
  149. package/skills/openspec/SKILL.md +346 -0
  150. package/skills/pi-extensions/SKILL.md +191 -0
  151. package/skills/pi-tui/SKILL.md +517 -0
  152. package/skills/python/SKILL.md +189 -0
  153. package/skills/rust/SKILL.md +268 -0
  154. package/skills/security/SKILL.md +206 -0
  155. package/skills/style/SKILL.md +264 -0
  156. package/skills/typescript/SKILL.md +225 -0
  157. package/skills/vault/SKILL.md +102 -0
  158. package/themes/alpharius-legacy.json +85 -0
  159. package/themes/alpharius.conf +59 -0
  160. package/themes/alpharius.json +88 -0
@@ -0,0 +1,752 @@
1
+ /**
2
+ * model-budget — Model tier + thinking level control
3
+ *
4
+ * Provides two orthogonal levers for cost/capability tuning:
5
+ * 1. Model tier: gloriana (deep) → victory (capable) → retribution (fast)
6
+ * 2. Thinking level: off → minimal → low → medium → high
7
+ *
8
+ * The agent can adjust both independently. Combined, these give fine-grained
9
+ * control: e.g., victory+high for moderate tasks that need careful reasoning,
10
+ * or gloriana+low for broad context understanding with minimal deliberation.
11
+ *
12
+ * Tools:
13
+ * set_model_tier — Switch model (gloriana/victory/retribution)
14
+ * set_thinking_level — Adjust extended thinking budget
15
+ *
16
+ * Commands:
17
+ * /gloriana, /victory, /retribution — Direct model switch
18
+ */
19
+
20
+ import { createHash } from "node:crypto";
21
+ import type { ExtensionAPI, ExtensionContext } from "@cwilson613/pi-coding-agent";
22
+ import type { ImageContent, Model, TextContent } from "@cwilson613/pi-ai";
23
+ import { DASHBOARD_UPDATE_EVENT, sharedState } from "./shared-state.ts";
24
+ import type { RecoveryEvent, RecoveryFailureClassification } from "./shared-state.ts";
25
+ import type { RecoveryAction, RecoveryCooldownSummary, RecoveryDashboardState, RecoveryTarget } from "./dashboard/types.ts";
26
+ import { tierConfig } from "./effort/tiers.ts";
27
+ import type { EffortLevel } from "./effort/types.ts";
28
+ import { clampThinkingLevel, classifyUpstreamFailure, getDefaultPolicy, getTierDisplayLabel, resolveTier, type CapabilityRuntimeState, type ModelTier, type RegistryModel, type UpstreamFailureClassification } from "./lib/model-routing.ts";
29
+ import { writeLastUsedModel } from "./lib/model-preferences.ts";
30
+ import { loadOperatorRuntimeState, readOperatorProfile, toCapabilityProfile, toCapabilityRuntimeState } from "./lib/operator-profile.ts";
31
+ import { buildFallbackGuidance, explainTierResolutionFailure, planRecoveryForModel, recordTransientFailureForModel, type RecoveryPlan } from "./lib/operator-fallback.ts";
32
+ import { switchToOfflineDriver } from "./offline-driver.ts";
33
+
34
+ /** Model tier ordering for effort cap comparison. */
35
+ export const TIER_ORDER: Record<string, number> = { local: 0, retribution: 1, victory: 2, gloriana: 3 };
36
+
37
+ /**
38
+ * Check whether an effort cap blocks a model tier switch.
39
+ *
40
+ * Derives the ceiling from capLevel (the level at which the cap was set),
41
+ * NOT from effort.driver (which reflects the current tier and changes
42
+ * when the operator switches tiers mid-session).
43
+ *
44
+ * If sharedState.effort is capped and the requested tier is higher than the
45
+ * cap ceiling's driver, returns { blocked: true, message: "..." }.
46
+ * Otherwise returns { blocked: false }.
47
+ *
48
+ * Exported for testing (extensions/effort/model-budget-cap.test.ts).
49
+ */
50
+ export function checkEffortCap(requestedTier: string): { blocked: boolean; message?: string } {
51
+ const effort = (sharedState as any).effort as
52
+ | { capped?: boolean; capLevel?: number; driver?: string; name?: string; level?: number }
53
+ | undefined;
54
+ if (!effort?.capped || effort.capLevel == null) return { blocked: false };
55
+
56
+ // Derive the ceiling driver from the capLevel, not the current tier's driver.
57
+ const capConfig = tierConfig(effort.capLevel as EffortLevel);
58
+ const capDriver = capConfig.driver;
59
+
60
+ const requestedOrder = TIER_ORDER[requestedTier] ?? -1;
61
+ const capOrder = TIER_ORDER[capDriver] ?? -1;
62
+
63
+ if (requestedOrder > capOrder) {
64
+ return {
65
+ blocked: true,
66
+ message:
67
+ `Effort cap active: ${capConfig.name} (level ${effort.capLevel}) limits driver to ${capDriver}. ` +
68
+ `Cannot upgrade to ${requestedTier}. Use /effort uncap to remove the ceiling.`,
69
+ };
70
+ }
71
+ return { blocked: false };
72
+ }
73
+
74
+ /** Tier icons for operator notifications */
75
+ const TIER_ICONS: Record<ModelTier, string> = {
76
+ local: "🤖",
77
+ retribution: "💨",
78
+ victory: "⚡",
79
+ gloriana: "🧠",
80
+ };
81
+
82
+ type TierName = ModelTier;
83
+
84
+ // Thinking levels ordered by cost/depth (xhigh excluded — OpenAI-only)
85
+ const THINKING_LEVELS = ["off", "minimal", "low", "medium", "high"] as const;
86
+ type ThinkingLevelName = (typeof THINKING_LEVELS)[number];
87
+
88
+ const THINKING_LABELS: Record<ThinkingLevelName, { icon: string; label: string }> = {
89
+ off: { icon: "⏭️", label: "no thinking" },
90
+ minimal: { icon: "💭", label: "minimal thinking" },
91
+ low: { icon: "💭", label: "low thinking" },
92
+ medium: { icon: "🤔", label: "medium thinking" },
93
+ high: { icon: "🧠", label: "deep thinking" },
94
+ };
95
+
96
+ const TIER_CAPABILITY_COPY: Record<TierName, string> = {
97
+ local: "on-device local execution",
98
+ retribution: "fast lightweight cloud tier",
99
+ victory: "balanced capability tier",
100
+ gloriana: "deep reasoning tier",
101
+ };
102
+
103
+ export function buildSetModelTierDescription(): string {
104
+ return (
105
+ "Switch the active capability tier based on task complexity. " +
106
+ "Omegon resolves the requested tier through the active provider routing policy, so the backing model may come from Anthropic, OpenAI, or local inference. " +
107
+ "Use 'local' for on-device work, 'retribution' for simple lookups and boilerplate, 'victory' for routine coding and execution, and 'gloriana' for deep reasoning and architecture."
108
+ );
109
+ }
110
+
111
+ export function buildTierCommandDescription(tier: TierName): string {
112
+ return `Switch to ${getTierDisplayLabel(tier)} — ${TIER_CAPABILITY_COPY[tier]} via provider-aware routing`;
113
+ }
114
+
115
+ function getResolverInputs(ctx: ExtensionContext) {
116
+ const policy = (sharedState as any).routingPolicy ?? getDefaultPolicy();
117
+ const profile = toCapabilityProfile(readOperatorProfile(ctx.cwd));
118
+ const runtimeState = toCapabilityRuntimeState(loadOperatorRuntimeState(ctx.cwd));
119
+ return { policy, profile, runtimeState };
120
+ }
121
+
122
+ function getAssistantErrorMessage(message: unknown): string | undefined {
123
+ if (!message || typeof message !== "object") return undefined;
124
+ const record = message as { role?: string; errorMessage?: unknown };
125
+ if (record.role !== "assistant" || typeof record.errorMessage !== "string" || !record.errorMessage.trim()) {
126
+ return undefined;
127
+ }
128
+ return record.errorMessage;
129
+ }
130
+
131
+ function summarizeErrorMessage(errorMessage: string): string {
132
+ return errorMessage.replace(/\s+/g, " ").trim().slice(0, 240);
133
+ }
134
+
135
+ function mapRecoveryFailureClassification(classification: UpstreamFailureClassification): {
136
+ classification: RecoveryFailureClassification;
137
+ retryable: boolean;
138
+ guidance: string;
139
+ } {
140
+ switch (classification.class) {
141
+ case "retryable-flake":
142
+ return {
143
+ classification: "transient_server_error",
144
+ retryable: true,
145
+ guidance: "Obvious upstream flakiness can retry once on the same provider/model before escalation.",
146
+ };
147
+ case "rate-limit":
148
+ case "backoff":
149
+ return {
150
+ classification: "rate_limited",
151
+ retryable: false,
152
+ guidance: "Rate limiting/backoff should cool down the failing route and prefer an alternate candidate instead of blind retry.",
153
+ };
154
+ case "auth":
155
+ return {
156
+ classification: "authentication_failed",
157
+ retryable: false,
158
+ guidance: "Authentication failed; refresh credentials or switch to a provider with a valid session.",
159
+ };
160
+ case "quota":
161
+ return {
162
+ classification: "quota_exhausted",
163
+ retryable: false,
164
+ guidance: "Quota exhaustion is not retryable; switch models/providers or restore quota before retrying.",
165
+ };
166
+ case "tool-output":
167
+ return {
168
+ classification: "malformed_output",
169
+ retryable: false,
170
+ guidance: "Malformed output should not use generic retry; adjust the prompt/schema or switch models explicitly.",
171
+ };
172
+ case "context-overflow":
173
+ return {
174
+ classification: "context_overflow",
175
+ retryable: false,
176
+ guidance: "Context overflow is handled separately; compact context or reduce prompt size before retrying.",
177
+ };
178
+ case "invalid-request":
179
+ return {
180
+ classification: "invalid_request",
181
+ retryable: false,
182
+ guidance: "The API rejected the request (e.g. image too large, malformed payload). Fix the request content before retrying.",
183
+ };
184
+ case "user-abort":
185
+ return {
186
+ classification: "unknown_upstream",
187
+ retryable: false,
188
+ guidance: "Operation was cancelled by the user; no recovery action needed.",
189
+ };
190
+ default:
191
+ return {
192
+ classification: "unknown_upstream",
193
+ retryable: false,
194
+ guidance: "Upstream failure was not safely classified for automatic retry; surface it and ask the operator to choose the next route.",
195
+ };
196
+ }
197
+ }
198
+
199
+ export function classifyRecoveryFailure(errorMessage: string): {
200
+ classification: RecoveryFailureClassification;
201
+ retryable: boolean;
202
+ guidance: string;
203
+ } {
204
+ return mapRecoveryFailureClassification(classifyUpstreamFailure(errorMessage));
205
+ }
206
+
207
+ function hashRecoveryRequestContent(content: string | (TextContent | ImageContent)[]): string {
208
+ return createHash("sha1").update(JSON.stringify(content)).digest("hex").slice(0, 12);
209
+ }
210
+
211
+ function getRetryLedgerKey(provider: string, model: string, requestFingerprint: string): string {
212
+ return `${provider}/${model}:${requestFingerprint}`;
213
+ }
214
+
215
+ function getLatestUserRecoveryRequest(ctx: ExtensionContext): {
216
+ content: string | (TextContent | ImageContent)[];
217
+ fingerprint: string;
218
+ } | undefined {
219
+ const entries = ctx.sessionManager.getEntries();
220
+ for (let index = entries.length - 1; index >= 0; index -= 1) {
221
+ const entry = entries[index];
222
+ if (entry.type !== "message" || entry.message.role !== "user") continue;
223
+ const content = entry.message.content as string | (TextContent | ImageContent)[];
224
+ return {
225
+ content,
226
+ fingerprint: hashRecoveryRequestContent(content),
227
+ };
228
+ }
229
+ return undefined;
230
+ }
231
+
232
+ export function piCoreAutoRetryLikelyHandles(errorMessage: string): boolean {
233
+ return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|terminated|retry delay/i.test(errorMessage);
234
+ }
235
+
236
+ export function buildRecoveryEvent(params: {
237
+ provider: string;
238
+ model: string;
239
+ turnIndex: number;
240
+ errorMessage: string;
241
+ retryCount: number;
242
+ guidance: string;
243
+ alternateCandidate?: { provider: string; id: string };
244
+ cooldownApplied?: boolean;
245
+ }): RecoveryEvent {
246
+ const classified = classifyRecoveryFailure(params.errorMessage);
247
+ const retryAttempted = classified.retryable && params.retryCount < 1;
248
+ return {
249
+ provider: params.provider,
250
+ model: params.model,
251
+ turnIndex: params.turnIndex,
252
+ classification: classified.classification,
253
+ originalErrorSummary: summarizeErrorMessage(params.errorMessage),
254
+ retryable: classified.retryable,
255
+ disposition: classified.classification === "context_overflow"
256
+ ? "handled_elsewhere"
257
+ : retryAttempted
258
+ ? "retry_same_model"
259
+ : params.cooldownApplied || classified.classification === "rate_limited"
260
+ ? "cooldown_and_failover"
261
+ : classified.retryable
262
+ ? "escalate"
263
+ : "guidance_only",
264
+ retryAttempted,
265
+ retryCount: retryAttempted ? params.retryCount + 1 : params.retryCount,
266
+ maxRetries: classified.retryable ? 1 : 0,
267
+ guidance: params.guidance || classified.guidance,
268
+ cooldownApplied: params.cooldownApplied,
269
+ alternateCandidate: params.alternateCandidate
270
+ ? { provider: params.alternateCandidate.provider, model: params.alternateCandidate.id }
271
+ : undefined,
272
+ timestamp: Date.now(),
273
+ };
274
+ }
275
+
276
+ function buildRecoveryCooldowns(runtimeState: CapabilityRuntimeState | undefined): RecoveryCooldownSummary[] | undefined {
277
+ if (!runtimeState) return undefined;
278
+ const cooldowns: RecoveryCooldownSummary[] = [];
279
+
280
+ for (const [provider, entry] of Object.entries(runtimeState.providerCooldowns ?? {})) {
281
+ if (!entry) continue;
282
+ cooldowns.push({
283
+ scope: "provider",
284
+ key: provider,
285
+ provider,
286
+ until: entry.until,
287
+ reason: entry.reason,
288
+ });
289
+ }
290
+
291
+ for (const [key, entry] of Object.entries(runtimeState.candidateCooldowns ?? {})) {
292
+ const [provider, modelId] = key.split("/");
293
+ cooldowns.push({
294
+ scope: "candidate",
295
+ key,
296
+ provider,
297
+ modelId,
298
+ until: entry.until,
299
+ reason: entry.reason,
300
+ });
301
+ }
302
+
303
+ return cooldowns.length > 0 ? cooldowns.sort((a, b) => a.until - b.until) : undefined;
304
+ }
305
+
306
+ function mapRecoveryAction(plan: RecoveryPlan, escalated = false): RecoveryAction {
307
+ if (escalated) return "escalate";
308
+ switch (plan.action) {
309
+ case "retry-same-model": return "retry";
310
+ case "switch-model": return "switch_candidate";
311
+ case "handoff-local": return "switch_offline";
312
+ case "handled-elsewhere": return "observe";
313
+ case "surface":
314
+ return plan.classification.cooldownProvider || plan.classification.cooldownCandidate ? "cooldown" : "observe";
315
+ default:
316
+ return "observe";
317
+ }
318
+ }
319
+
320
+ function buildRecoverySummary(plan: RecoveryPlan, target: RecoveryTarget | undefined, escalated = false): string {
321
+ if (escalated) {
322
+ return `Escalated ${plan.classification.summary} after recovery could not switch away from the failing route.`;
323
+ }
324
+ switch (plan.action) {
325
+ case "retry-same-model":
326
+ return `Retrying once on the same model after ${plan.classification.summary}.`;
327
+ case "switch-model":
328
+ return target?.modelId
329
+ ? `Switching recovery to ${target.provider}/${target.modelId} after ${plan.classification.summary}.`
330
+ : `Switching to an alternate candidate after ${plan.classification.summary}.`;
331
+ case "handoff-local":
332
+ return target?.label
333
+ ? `Switching recovery to local driver ${target.label} after ${plan.classification.summary}.`
334
+ : `Switching recovery to the local driver after ${plan.classification.summary}.`;
335
+ case "handled-elsewhere":
336
+ return `Observed ${plan.classification.summary}; recovery is handled by explicit compaction/context management logic.`;
337
+ default:
338
+ return `Observed ${plan.classification.summary}; ${plan.reason}`;
339
+ }
340
+ }
341
+
342
+ export function buildRecoveryDashboardState(params: {
343
+ recoveryEvent: RecoveryEvent;
344
+ plan: RecoveryPlan;
345
+ runtimeState?: CapabilityRuntimeState;
346
+ target?: RecoveryTarget;
347
+ escalated?: boolean;
348
+ }): RecoveryDashboardState {
349
+ return {
350
+ provider: params.recoveryEvent.provider,
351
+ modelId: params.recoveryEvent.model,
352
+ classification: params.recoveryEvent.classification,
353
+ summary: buildRecoverySummary(params.plan, params.target, params.escalated ?? false),
354
+ action: mapRecoveryAction(params.plan, params.escalated ?? false),
355
+ retryCount: params.recoveryEvent.retryCount,
356
+ maxRetries: params.recoveryEvent.maxRetries,
357
+ attemptId: `${params.recoveryEvent.turnIndex}:${params.recoveryEvent.provider}/${params.recoveryEvent.model}`,
358
+ timestamp: params.recoveryEvent.timestamp,
359
+ escalated: params.escalated,
360
+ target: params.target,
361
+ cooldowns: buildRecoveryCooldowns(params.runtimeState),
362
+ };
363
+ }
364
+
365
+ function buildRecoveryNotice(recoveryEvent: RecoveryEvent, dashboardState: RecoveryDashboardState): string {
366
+ const target = dashboardState.target?.modelId
367
+ ? `${dashboardState.target.provider}/${dashboardState.target.modelId}`
368
+ : dashboardState.target?.provider;
369
+ const retry = recoveryEvent.maxRetries > 0 ? `retry ${recoveryEvent.retryCount}/${recoveryEvent.maxRetries}` : undefined;
370
+ return [
371
+ `Recovery observed ${recoveryEvent.classification} for ${recoveryEvent.provider}/${recoveryEvent.model}.`,
372
+ dashboardState.summary,
373
+ retry,
374
+ target ? `target ${target}` : undefined,
375
+ ].filter(Boolean).join(" ");
376
+ }
377
+
378
+ export function shouldUseExtensionRetryFallback(errorMessage: string, retryAttempted: boolean): boolean {
379
+ return retryAttempted && !piCoreAutoRetryLikelyHandles(errorMessage);
380
+ }
381
+
382
+ function scheduleExtensionRetry(
383
+ pi: ExtensionAPI,
384
+ request: { content: string | (TextContent | ImageContent)[]; fingerprint: string },
385
+ ): void {
386
+ setTimeout(() => {
387
+ try {
388
+ pi.sendUserMessage(request.content);
389
+ } catch {
390
+ // If the retry prompt itself fails to queue, the original recovery notice remains in-session.
391
+ }
392
+ }, 0);
393
+ }
394
+
395
+ async function applyRecoveryPlan(
396
+ plan: RecoveryPlan,
397
+ pi: ExtensionAPI,
398
+ ctx: ExtensionContext,
399
+ ): Promise<{ target?: RecoveryTarget; escalated?: boolean }> {
400
+ if (plan.action === "switch-model" && plan.alternateCandidate) {
401
+ const targetModel = ctx.modelRegistry.find(plan.alternateCandidate.provider, plan.alternateCandidate.id);
402
+ if (!targetModel) return { escalated: true };
403
+ const success = await pi.setModel(targetModel as Model<any>);
404
+ return success
405
+ ? { target: { provider: plan.alternateCandidate.provider, modelId: plan.alternateCandidate.id, label: `${plan.alternateCandidate.provider}/${plan.alternateCandidate.id}` } }
406
+ : { escalated: true };
407
+ }
408
+
409
+ if (plan.action === "handoff-local") {
410
+ const offline = await switchToOfflineDriver(pi, ctx as any, {
411
+ preferredModel: plan.alternateCandidate?.id,
412
+ automatic: true,
413
+ });
414
+ return offline.success
415
+ ? { target: { provider: offline.provider, modelId: offline.modelId, label: offline.label } }
416
+ : { escalated: true };
417
+ }
418
+
419
+ return {};
420
+ }
421
+
422
+ async function switchTo(tier: TierName, pi: ExtensionAPI, ctx: ExtensionContext): Promise<RegistryModel | null> {
423
+ const all = ctx.modelRegistry.getAll() as unknown as RegistryModel[];
424
+ const { policy, profile, runtimeState } = getResolverInputs(ctx);
425
+ const resolved = resolveTier(tier, all, policy, runtimeState, profile);
426
+ if (!resolved) return null;
427
+ const model = all.find((m) => m.id === resolved.modelId);
428
+ if (!model) return null;
429
+ const success = await pi.setModel(model as unknown as Model<any>);
430
+ if (success) {
431
+ writeLastUsedModel(ctx.cwd, { provider: model.provider, modelId: model.id });
432
+ const currentThinking = pi.getThinkingLevel() as ThinkingLevelName;
433
+ const clampedThinking = clampThinkingLevel(currentThinking, resolved.maxThinking ?? "high");
434
+ if (clampedThinking !== currentThinking) {
435
+ pi.setThinkingLevel(clampedThinking as any);
436
+ }
437
+ return model;
438
+ }
439
+ return null;
440
+ }
441
+
442
+ function currentTierName(ctx: ExtensionContext): TierName | null {
443
+ const model = ctx.model;
444
+ if (!model) return null;
445
+ // Resolve the current model against the registry using the shared resolver
446
+ const all = ctx.modelRegistry.getAll() as unknown as RegistryModel[];
447
+ const { policy, profile, runtimeState } = getResolverInputs(ctx);
448
+ for (const tier of ["gloriana", "victory", "retribution", "local"] as TierName[]) {
449
+ const resolved = resolveTier(tier, all, policy, runtimeState, profile);
450
+ if (resolved?.modelId === model.id) return tier;
451
+ }
452
+ return null;
453
+ }
454
+
455
+ export default function (pi: ExtensionAPI) {
456
+ // session_start model selection is handled by the effort extension.
457
+ // model-budget only provides the set_model_tier / set_thinking_level tools.
458
+
459
+ pi.on("turn_end", async (event, ctx) => {
460
+ const errorMessage = getAssistantErrorMessage(event.message);
461
+ if (!errorMessage) {
462
+ if (sharedState.recoveryRetryCounts && Object.keys(sharedState.recoveryRetryCounts).length > 0) {
463
+ sharedState.recoveryRetryCounts = {};
464
+ }
465
+ // Clear stale recovery state on a successful turn so the banner doesn't
466
+ // linger indefinitely in the compact footer.
467
+ if (sharedState.recovery || sharedState.latestRecoveryEvent) {
468
+ sharedState.recovery = undefined;
469
+ sharedState.latestRecoveryEvent = undefined;
470
+ pi.events.emit(DASHBOARD_UPDATE_EVENT, { source: "model-budget", recovery: undefined });
471
+ }
472
+ return;
473
+ }
474
+
475
+ // User-initiated aborts (Esc / SIGINT / AbortSignal) must never enter the
476
+ // upstream recovery pipeline. They are not API failures.
477
+ const abortClassification = classifyUpstreamFailure(errorMessage);
478
+ if (abortClassification.class === "user-abort") {
479
+ // Clear any stale recovery state so the dashboard doesn't linger on the
480
+ // previous recovery notice.
481
+ sharedState.recovery = undefined;
482
+ sharedState.latestRecoveryEvent = undefined;
483
+ pi.events.emit(DASHBOARD_UPDATE_EVENT, { source: "model-budget", recovery: undefined });
484
+ return;
485
+ }
486
+
487
+ if (!ctx.model) return;
488
+
489
+ const provider = ctx.model.provider;
490
+ const recoveryRequest = getLatestUserRecoveryRequest(ctx);
491
+ const ledgerKey = getRetryLedgerKey(
492
+ provider,
493
+ ctx.model.id,
494
+ recoveryRequest?.fingerprint ?? `turn-${event.turnIndex}`,
495
+ );
496
+ const retryCounts = sharedState.recoveryRetryCounts ?? {};
497
+ const priorRetryCount = retryCounts[ledgerKey] ?? 0;
498
+
499
+ const persistedRuntimeState = recordTransientFailureForModel(ctx.cwd, ctx.model, errorMessage);
500
+ const { policy, profile } = getResolverInputs(ctx);
501
+ const models = ctx.modelRegistry.getAll() as unknown as RegistryModel[];
502
+ const runtimeState = persistedRuntimeState ?? toCapabilityRuntimeState(loadOperatorRuntimeState(ctx.cwd));
503
+ const plan = planRecoveryForModel(ctx.model, errorMessage, models, policy, profile, runtimeState ?? {}, Date.now());
504
+ const guidance = buildFallbackGuidance(ctx.model, models, policy, profile, runtimeState ?? {}, Date.now());
505
+ const applied = await applyRecoveryPlan(plan, pi, ctx);
506
+ const classified = mapRecoveryFailureClassification(classifyUpstreamFailure(errorMessage));
507
+ const recoveryEvent = buildRecoveryEvent({
508
+ provider,
509
+ model: ctx.model.id,
510
+ turnIndex: event.turnIndex,
511
+ errorMessage,
512
+ retryCount: priorRetryCount,
513
+ guidance: plan.reason || guidance?.reason || classified.guidance,
514
+ alternateCandidate: plan.alternateCandidate,
515
+ cooldownApplied: Boolean(persistedRuntimeState),
516
+ });
517
+ const dashboardState = buildRecoveryDashboardState({
518
+ recoveryEvent,
519
+ plan,
520
+ runtimeState,
521
+ target: applied.target,
522
+ escalated: applied.escalated,
523
+ });
524
+ const useExtensionRetryFallback = shouldUseExtensionRetryFallback(errorMessage, recoveryEvent.retryAttempted);
525
+
526
+ sharedState.latestRecoveryEvent = recoveryEvent;
527
+ sharedState.recovery = dashboardState;
528
+ sharedState.recoveryRetryCounts = {
529
+ ...retryCounts,
530
+ [ledgerKey]: recoveryEvent.retryCount,
531
+ };
532
+ pi.sendMessage({
533
+ customType: "recovery-event",
534
+ content: buildRecoveryNotice(recoveryEvent, dashboardState),
535
+ display: true,
536
+ details: {
537
+ recoveryEvent,
538
+ recovery: dashboardState,
539
+ plan: {
540
+ action: plan.action,
541
+ sameModelRetry: plan.sameModelRetry,
542
+ reason: plan.reason,
543
+ classification: plan.classification.class,
544
+ role: plan.role,
545
+ alternateCandidate: plan.alternateCandidate,
546
+ retryStrategy: useExtensionRetryFallback ? "extension-sendUserMessage" : "core-auto-retry",
547
+ },
548
+ },
549
+ }, { triggerTurn: false });
550
+ pi.events.emit(DASHBOARD_UPDATE_EVENT, { source: "model-budget", recoveryEvent, recovery: dashboardState });
551
+
552
+ if (useExtensionRetryFallback && recoveryRequest) {
553
+ scheduleExtensionRetry(pi, recoveryRequest);
554
+ }
555
+
556
+ if (!ctx.hasUI) return;
557
+
558
+ if (dashboardState.action === "retry") {
559
+ ctx.ui.notify(dashboardState.summary, "warning");
560
+ return;
561
+ }
562
+
563
+ if (dashboardState.action === "switch_candidate" || dashboardState.action === "switch_offline") {
564
+ ctx.ui.notify(dashboardState.summary, applied.escalated ? "warning" : "info");
565
+ return;
566
+ }
567
+
568
+ const level = recoveryEvent.disposition === "handled_elsewhere"
569
+ ? "info"
570
+ : guidance?.requiresConfirmation || applied.escalated
571
+ ? "warning"
572
+ : recoveryEvent.retryable
573
+ ? "warning"
574
+ : "error";
575
+ ctx.ui.notify(dashboardState.summary, level);
576
+ });
577
+
578
+ const modelTierParameters = {
579
+ type: "object",
580
+ properties: {
581
+ tier: {
582
+ type: "string",
583
+ enum: ["local", "retribution", "victory", "gloriana"],
584
+ description: "Target model tier",
585
+ },
586
+ reason: {
587
+ type: "string",
588
+ description: "Brief explanation for the tier change",
589
+ },
590
+ },
591
+ required: ["tier", "reason"],
592
+ additionalProperties: false,
593
+ } as const;
594
+
595
+ const thinkingLevelParameters = {
596
+ type: "object",
597
+ properties: {
598
+ level: {
599
+ type: "string",
600
+ enum: ["off", "minimal", "low", "medium", "high"],
601
+ description: "Thinking level — higher = more reasoning tokens, slower, more expensive",
602
+ },
603
+ reason: {
604
+ type: "string",
605
+ description: "Brief explanation for the thinking level change",
606
+ },
607
+ },
608
+ required: ["level", "reason"],
609
+ additionalProperties: false,
610
+ } as const;
611
+
612
+ // --- Model Tier Tool ---
613
+ pi.registerTool({
614
+ name: "set_model_tier",
615
+ label: "Set Model Tier",
616
+ description: buildSetModelTierDescription(),
617
+ promptSnippet: "Switch capability tier (local/retribution/victory/gloriana) through provider-aware routing",
618
+ promptGuidelines: [
619
+ "Downgrade to victory for routine file edits, command execution, and cleanup tasks",
620
+ "Upgrade to gloriana when encountering architecture decisions, complex debugging, or multi-step planning",
621
+ "Use retribution for simple lookups, formatting, and boilerplate generation",
622
+ ],
623
+ parameters: modelTierParameters as any,
624
+ execute: async (
625
+ _toolCallId,
626
+ params: { tier: string; reason: string },
627
+ _signal,
628
+ _onUpdate,
629
+ ctx,
630
+ ) => {
631
+ const tier = params.tier as TierName;
632
+ const icon = TIER_ICONS[tier];
633
+ const displayLabel = getTierDisplayLabel(tier);
634
+
635
+ // Enforce effort cap — block upgrades past the ceiling
636
+ const capCheck = checkEffortCap(tier);
637
+ if (capCheck.blocked) {
638
+ return {
639
+ content: [{ type: "text" as const, text: capCheck.message! }],
640
+ details: undefined,
641
+ };
642
+ }
643
+
644
+ const model = await switchTo(tier, pi, ctx);
645
+ if (model) {
646
+ const thinking = pi.getThinkingLevel();
647
+ const target = `${model.provider}/${model.id}`;
648
+ ctx.ui.notify(`${icon} → ${displayLabel} [${tier}] → ${target} (thinking: ${thinking}): ${params.reason}`, "info");
649
+ return {
650
+ content: [
651
+ {
652
+ type: "text" as const,
653
+ text: `Switched to ${displayLabel} [${tier}] via ${target}, thinking: ${thinking}. ${params.reason}`,
654
+ },
655
+ ],
656
+ details: undefined,
657
+ };
658
+ }
659
+ const { policy, profile, runtimeState } = getResolverInputs(ctx);
660
+ const failure = explainTierResolutionFailure(
661
+ tier,
662
+ ctx.modelRegistry.getAll() as unknown as RegistryModel[],
663
+ policy,
664
+ profile,
665
+ runtimeState,
666
+ ) ?? `Failed to switch to ${displayLabel} [${tier}] — no matching model found or no API key`;
667
+ return {
668
+ content: [
669
+ {
670
+ type: "text" as const,
671
+ text: failure,
672
+ },
673
+ ],
674
+ details: undefined,
675
+ };
676
+ },
677
+ });
678
+
679
+ // --- Thinking Level Tool ---
680
+ pi.registerTool({
681
+ name: "set_thinking_level",
682
+ label: "Set Thinking Level",
683
+ description:
684
+ "Adjust the extended thinking budget independently of model tier. " +
685
+ "Higher levels allocate more tokens for internal reasoning before responding. " +
686
+ "Use 'high' for complex multi-step problems, debugging, or architecture. " +
687
+ "Use 'medium' (default) for general tasks. " +
688
+ "Use 'low' or 'minimal' for straightforward execution where speed matters. " +
689
+ "Use 'off' to disable extended thinking entirely (fastest, cheapest). " +
690
+ "Thinking level and model tier are orthogonal — adjust both for fine-grained control.",
691
+ promptSnippet: "Adjust extended thinking budget (off/minimal/low/medium/high)",
692
+ promptGuidelines: [
693
+ "Reduce thinking for mechanical tasks: file reads, grep, simple edits, formatting",
694
+ "Increase thinking for: debugging, architecture decisions, complex refactors, multi-file changes",
695
+ "Combine with model tier: victory+high is cheaper than gloriana+medium for moderate reasoning tasks",
696
+ ],
697
+ parameters: thinkingLevelParameters as any,
698
+ execute: async (
699
+ _toolCallId,
700
+ params: { level: string; reason: string },
701
+ _signal,
702
+ _onUpdate,
703
+ ctx,
704
+ ) => {
705
+ const previous = pi.getThinkingLevel();
706
+ pi.setThinkingLevel(params.level as any);
707
+ const level = params.level as ThinkingLevelName;
708
+ const info = THINKING_LABELS[level];
709
+ const tier = currentTierName(ctx) ?? "unknown";
710
+ ctx.ui.notify(`${info.icon} thinking: ${previous} → ${level} (model: ${tier}): ${params.reason}`, "info");
711
+ return {
712
+ content: [
713
+ {
714
+ type: "text" as const,
715
+ text: `Thinking: ${previous} → ${level} (${info.label}), model: ${tier}. ${params.reason}`,
716
+ },
717
+ ],
718
+ details: undefined,
719
+ };
720
+ },
721
+ });
722
+
723
+ // --- Manual commands for direct control ---
724
+ const COMMAND_TIERS: ModelTier[] = ["local", "retribution", "victory", "gloriana"];
725
+ for (const tier of COMMAND_TIERS) {
726
+ const icon = TIER_ICONS[tier];
727
+ const displayLabel = getTierDisplayLabel(tier);
728
+ pi.registerCommand(tier, {
729
+ description: `${buildTierCommandDescription(tier)} (${icon})`,
730
+ handler: async (_args, ctx) => {
731
+ // Enforce effort cap — same check as the tool
732
+ const capCheck = checkEffortCap(tier);
733
+ if (capCheck.blocked) {
734
+ ctx.ui.notify(`⛔ ${capCheck.message}`, "warning");
735
+ return;
736
+ }
737
+ const model = await switchTo(tier, pi, ctx);
738
+ if (!model) {
739
+ const { policy, profile, runtimeState } = getResolverInputs(ctx);
740
+ const failure = explainTierResolutionFailure(
741
+ tier,
742
+ ctx.modelRegistry.getAll() as unknown as RegistryModel[],
743
+ policy,
744
+ profile,
745
+ runtimeState,
746
+ );
747
+ ctx.ui.notify(failure ?? `Failed to switch to ${displayLabel} [${tier}]`, "error");
748
+ }
749
+ },
750
+ });
751
+ }
752
+ }