@oh-my-pi/pi-coding-agent 14.1.0 → 14.1.2

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 (82) hide show
  1. package/CHANGELOG.md +79 -0
  2. package/package.json +8 -8
  3. package/src/async/job-manager.ts +43 -10
  4. package/src/commit/agentic/tools/analyze-file.ts +1 -2
  5. package/src/config/mcp-schema.json +1 -1
  6. package/src/config/model-equivalence.ts +1 -0
  7. package/src/config/model-registry.ts +63 -34
  8. package/src/config/model-resolver.ts +111 -15
  9. package/src/config/settings-schema.ts +4 -3
  10. package/src/config/settings.ts +1 -1
  11. package/src/cursor.ts +64 -23
  12. package/src/edit/index.ts +254 -89
  13. package/src/edit/modes/chunk.ts +336 -57
  14. package/src/edit/modes/hashline.ts +51 -26
  15. package/src/edit/modes/patch.ts +16 -10
  16. package/src/edit/modes/replace.ts +15 -7
  17. package/src/edit/renderer.ts +248 -94
  18. package/src/export/html/template.generated.ts +1 -1
  19. package/src/export/html/template.js +6 -4
  20. package/src/extensibility/custom-tools/types.ts +0 -3
  21. package/src/extensibility/extensions/loader.ts +16 -0
  22. package/src/extensibility/extensions/runner.ts +2 -7
  23. package/src/extensibility/extensions/types.ts +8 -4
  24. package/src/internal-urls/docs-index.generated.ts +3 -3
  25. package/src/ipy/executor.ts +447 -52
  26. package/src/ipy/kernel.ts +39 -13
  27. package/src/lsp/client.ts +54 -0
  28. package/src/lsp/index.ts +8 -0
  29. package/src/lsp/types.ts +6 -0
  30. package/src/main.ts +0 -1
  31. package/src/modes/acp/acp-agent.ts +4 -1
  32. package/src/modes/components/bash-execution.ts +16 -4
  33. package/src/modes/components/status-line/presets.ts +17 -6
  34. package/src/modes/components/status-line/segments.ts +15 -0
  35. package/src/modes/components/status-line-segment-editor.ts +1 -0
  36. package/src/modes/components/status-line.ts +7 -1
  37. package/src/modes/components/tool-execution.ts +145 -75
  38. package/src/modes/controllers/command-controller.ts +24 -1
  39. package/src/modes/controllers/event-controller.ts +4 -1
  40. package/src/modes/controllers/extension-ui-controller.ts +28 -5
  41. package/src/modes/controllers/input-controller.ts +9 -3
  42. package/src/modes/controllers/selector-controller.ts +4 -1
  43. package/src/modes/interactive-mode.ts +19 -3
  44. package/src/modes/print-mode.ts +13 -4
  45. package/src/modes/prompt-action-autocomplete.ts +3 -5
  46. package/src/modes/rpc/rpc-mode.ts +8 -2
  47. package/src/modes/shared.ts +2 -2
  48. package/src/modes/types.ts +1 -0
  49. package/src/modes/utils/ui-helpers.ts +1 -0
  50. package/src/prompts/tools/bash.md +2 -2
  51. package/src/prompts/tools/chunk-edit.md +191 -163
  52. package/src/prompts/tools/hashline.md +11 -11
  53. package/src/prompts/tools/patch.md +10 -5
  54. package/src/prompts/tools/{await.md → poll.md} +1 -1
  55. package/src/prompts/tools/read-chunk.md +3 -3
  56. package/src/prompts/tools/task.md +2 -2
  57. package/src/prompts/tools/vim.md +98 -0
  58. package/src/sdk.ts +754 -724
  59. package/src/session/agent-session.ts +164 -34
  60. package/src/session/session-manager.ts +50 -4
  61. package/src/slash-commands/builtin-registry.ts +17 -0
  62. package/src/task/executor.ts +4 -4
  63. package/src/task/index.ts +3 -5
  64. package/src/task/types.ts +2 -2
  65. package/src/tools/bash.ts +26 -8
  66. package/src/tools/find.ts +5 -2
  67. package/src/tools/grep.ts +77 -8
  68. package/src/tools/index.ts +48 -19
  69. package/src/tools/{await-tool.ts → poll-tool.ts} +36 -30
  70. package/src/tools/python.ts +293 -278
  71. package/src/tools/submit-result.ts +5 -2
  72. package/src/tools/todo-write.ts +8 -2
  73. package/src/tools/vim.ts +966 -0
  74. package/src/utils/edit-mode.ts +2 -1
  75. package/src/utils/session-color.ts +55 -0
  76. package/src/utils/title-generator.ts +15 -6
  77. package/src/vim/buffer.ts +309 -0
  78. package/src/vim/commands.ts +382 -0
  79. package/src/vim/engine.ts +2426 -0
  80. package/src/vim/parser.ts +151 -0
  81. package/src/vim/render.ts +252 -0
  82. package/src/vim/types.ts +197 -0
package/src/tools/bash.ts CHANGED
@@ -73,6 +73,7 @@ export interface BashToolInput {
73
73
 
74
74
  export interface BashToolDetails {
75
75
  meta?: OutputMeta;
76
+ timeoutSeconds?: number;
76
77
  async?: {
77
78
  state: "running" | "completed" | "failed";
78
79
  jobId: string;
@@ -290,14 +291,20 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
290
291
  tailLines?: number,
291
292
  ): AgentToolResult<BashToolDetails> {
292
293
  const outputText = this.#formatResultOutput(result, headLines, tailLines);
293
- const details: BashToolDetails = {};
294
+ const details: BashToolDetails = { timeoutSeconds: timeoutSec };
294
295
  const resultBuilder = toolResult(details).text(outputText).truncationFromSummary(result, { direction: "tail" });
295
296
  this.#buildResultText(result, timeoutSec, outputText);
296
297
  return resultBuilder.done();
297
298
  }
298
299
 
299
- #buildBackgroundStartResult(jobId: string, label: string, previewText: string): AgentToolResult<BashToolDetails> {
300
+ #buildBackgroundStartResult(
301
+ jobId: string,
302
+ label: string,
303
+ previewText: string,
304
+ timeoutSec: number,
305
+ ): AgentToolResult<BashToolDetails> {
300
306
  const details: BashToolDetails = {
307
+ timeoutSeconds: timeoutSec,
301
308
  async: { state: "running", jobId, type: "bash" },
302
309
  };
303
310
  const lines: string[] = [];
@@ -307,7 +314,7 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
307
314
  }
308
315
  lines.push(`Background job ${jobId} started: ${label}`);
309
316
  lines.push("Result will be delivered automatically when complete.");
310
- lines.push(`Use \`await\`, \`read jobs://${jobId}\`, or \`cancel_job\` if needed.`);
317
+ lines.push(`Use \`poll\`, \`read jobs://${jobId}\`, or \`cancel_job\` if needed.`);
311
318
  return {
312
319
  content: [{ type: "text", text: lines.join("\n") }],
313
320
  details,
@@ -430,6 +437,12 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
430
437
  }
431
438
  }
432
439
 
440
+ #resolveAutoBackgroundWaitMs(timeoutMs: number): number {
441
+ if (this.#autoBackgroundThresholdMs <= 0) return 0;
442
+ const timeoutBufferMs = 1_000;
443
+ return Math.max(0, Math.min(this.#autoBackgroundThresholdMs, timeoutMs - timeoutBufferMs));
444
+ }
445
+
433
446
  async execute(
434
447
  _toolCallId: string,
435
448
  {
@@ -536,10 +549,12 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
536
549
  onUpdate,
537
550
  startBackgrounded: true,
538
551
  });
539
- return this.#buildBackgroundStartResult(job.jobId, job.label, "");
552
+ return this.#buildBackgroundStartResult(job.jobId, job.label, "", timeoutSec);
540
553
  }
541
554
 
542
555
  if (this.#autoBackgroundEnabled && !pty && this.session.asyncJobManager) {
556
+ const autoBackgroundWaitMs = this.#resolveAutoBackgroundWaitMs(timeoutMs);
557
+ const startBackgrounded = autoBackgroundWaitMs === 0;
543
558
  const job = this.#startManagedBashJob({
544
559
  command,
545
560
  commandCwd,
@@ -549,9 +564,12 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
549
564
  tailLines,
550
565
  resolvedEnv,
551
566
  onUpdate,
552
- startBackgrounded: false,
567
+ startBackgrounded,
553
568
  });
554
- const waitResult = await this.#waitForManagedBashJob(job, this.#autoBackgroundThresholdMs, signal);
569
+ if (startBackgrounded) {
570
+ return this.#buildBackgroundStartResult(job.jobId, job.label, "", timeoutSec);
571
+ }
572
+ const waitResult = await this.#waitForManagedBashJob(job, autoBackgroundWaitMs, signal);
555
573
  if (waitResult.kind === "completed") {
556
574
  this.session.asyncJobManager.acknowledgeDeliveries([job.jobId]);
557
575
  return waitResult.result;
@@ -566,7 +584,7 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
566
584
  throw new ToolAbortError(job.getLatestText() || "Command aborted");
567
585
  }
568
586
  job.setBackgrounded(true);
569
- return this.#buildBackgroundStartResult(job.jobId, job.label, job.getLatestText());
587
+ return this.#buildBackgroundStartResult(job.jobId, job.label, job.getLatestText(), timeoutSec);
570
588
  }
571
589
 
572
590
  // Track output for streaming updates (tail only)
@@ -689,7 +707,7 @@ export const bashToolRenderer = {
689
707
  const showingFullOutput = expanded && renderContext?.isFullOutput === true;
690
708
 
691
709
  // Build truncation warning
692
- const timeoutSeconds = renderContext?.timeout;
710
+ const timeoutSeconds = details?.timeoutSeconds ?? renderContext?.timeout;
693
711
  const timeoutLine =
694
712
  typeof timeoutSeconds === "number"
695
713
  ? uiTheme.fg(
package/src/tools/find.ts CHANGED
@@ -244,17 +244,20 @@ export class FindTool implements AgentTool<typeof findSchema, FindToolDetails> {
244
244
  maxResults: effectiveLimit,
245
245
  sortByMtime: true,
246
246
  gitignore: useGitignore,
247
+ signal: combinedSignal,
247
248
  },
248
249
  onMatch,
249
- this.session.searchDb,
250
250
  ),
251
251
  );
252
252
 
253
253
  try {
254
254
  let result = await doGlob(true);
255
- if (result.matches.length === 0) {
255
+ if (result.matches.length === 0 && !timeoutSignal.aborted) {
256
256
  result = await doGlob(false);
257
257
  }
258
+ // Sort by mtime descending (most recent first) in JS instead of native.
259
+ // This allows native glob to early-terminate at maxResults.
260
+ result.matches.sort((a, b) => (b.mtime ?? 0) - (a.mtime ?? 0));
258
261
  matches = result.matches;
259
262
  } catch (error) {
260
263
  if (error instanceof Error && error.name === "AbortError") {
package/src/tools/grep.ts CHANGED
@@ -7,7 +7,7 @@ import { Text } from "@oh-my-pi/pi-tui";
7
7
  import { prompt, untilAborted } from "@oh-my-pi/pi-utils";
8
8
  import { type Static, Type } from "@sinclair/typebox";
9
9
  import { computeLineHash } from "../edit/line-hash";
10
- import { formatChunkedGrepLine } from "../edit/modes/chunk";
10
+ import { type ChunkedGrepMatch, describeChunkedGrepMatch } from "../edit/modes/chunk";
11
11
  import type { RenderResultOptions } from "../extensibility/custom-tools/types";
12
12
  import { getLanguageFromPath, type Theme } from "../modes/theme/theme";
13
13
  import grepDescription from "../prompts/tools/grep.md" with { type: "text" };
@@ -162,7 +162,8 @@ export class GrepTool implements AgentTool<typeof grepSchema, GrepToolDetails> {
162
162
  const stat = await Bun.file(searchPath).stat();
163
163
  isDirectory = stat.isDirectory();
164
164
  } catch {
165
- throw new ToolError(`Path not found: ${scopePath}`);
165
+ const hint = scopePath.includes(",") ? ` (comma-separated paths must each exist relative to cwd)` : "";
166
+ throw new ToolError(`Path not found: ${scopePath}${hint}`);
166
167
  }
167
168
 
168
169
  const effectiveOutputMode = GrepOutputMode.Content;
@@ -191,7 +192,6 @@ export class GrepTool implements AgentTool<typeof grepSchema, GrepToolDetails> {
191
192
  mode: effectiveOutputMode,
192
193
  },
193
194
  undefined,
194
- this.session.searchDb,
195
195
  );
196
196
  } catch (err) {
197
197
  if (err instanceof Error && err.message.startsWith("regex parse error")) {
@@ -274,13 +274,11 @@ export class GrepTool implements AgentTool<typeof grepSchema, GrepToolDetails> {
274
274
  matchesByFile.get(relativePath)!.push(match);
275
275
  }
276
276
  if (chunkMode) {
277
- const annotatedLines = await Promise.all(
277
+ const annotatedMatches = await Promise.all(
278
278
  selectedMatches.map(match => {
279
279
  const relativePath = match.path.startsWith("/") ? match.path.slice(1) : match.path;
280
280
  const absoluteFilePath = isDirectory ? path.join(searchPath, relativePath) : searchPath;
281
- const displayPath = formatPath(match.path);
282
- fileMatchCounts.set(displayPath, (fileMatchCounts.get(displayPath) ?? 0) + 1);
283
- return formatChunkedGrepLine({
281
+ return describeChunkedGrepMatch({
284
282
  filePath: absoluteFilePath,
285
283
  lineNumber: match.lineNumber,
286
284
  line: match.line,
@@ -289,7 +287,78 @@ export class GrepTool implements AgentTool<typeof grepSchema, GrepToolDetails> {
289
287
  });
290
288
  }),
291
289
  );
292
- const rawOutput = annotatedLines.join("\n");
290
+ const chunkMatchesByFile = new Map<string, ChunkedGrepMatch[]>();
291
+ for (const match of annotatedMatches) {
292
+ recordFile(match.displayPath);
293
+ if (!chunkMatchesByFile.has(match.displayPath)) {
294
+ chunkMatchesByFile.set(match.displayPath, []);
295
+ }
296
+ chunkMatchesByFile.get(match.displayPath)!.push(match);
297
+ }
298
+ const renderChunkedMatchesForFile = (relativePath: string) => {
299
+ const fileMatches = chunkMatchesByFile.get(relativePath) ?? [];
300
+ if (fileMatches.length === 0) {
301
+ return;
302
+ }
303
+ const lineWidth = fileMatches[0]?.fileLineCount.toString().length ?? 1;
304
+ const matchesByChunk = new Map<string, ChunkedGrepMatch[]>();
305
+ for (const match of fileMatches) {
306
+ const chunkKey = match.chunkPath ?? "";
307
+ if (!matchesByChunk.has(chunkKey)) {
308
+ matchesByChunk.set(chunkKey, []);
309
+ }
310
+ matchesByChunk.get(chunkKey)!.push(match);
311
+ }
312
+ for (const [chunkPath, chunkMatches] of matchesByChunk) {
313
+ if (chunkPath) {
314
+ const chunkChecksum = chunkMatches[0]?.chunkChecksum;
315
+ const dashes = "-".repeat(chunkPath.split(".").length - 1);
316
+ const anchor = chunkChecksum
317
+ ? `${dashes}@${chunkPath}#${chunkChecksum}`
318
+ : `${dashes}@${chunkPath}`;
319
+ outputLines.push(anchor);
320
+ }
321
+ for (const match of chunkMatches) {
322
+ outputLines.push(` ${match.lineNumber.toString().padStart(lineWidth, " ")} |${match.line}`);
323
+ fileMatchCounts.set(relativePath, (fileMatchCounts.get(relativePath) ?? 0) + 1);
324
+ }
325
+ }
326
+ };
327
+ if (isDirectory) {
328
+ const filesByDirectory = new Map<string, string[]>();
329
+ for (const relativePath of fileList) {
330
+ const directory = path.dirname(relativePath).replace(/\\/g, "/");
331
+ if (!filesByDirectory.has(directory)) {
332
+ filesByDirectory.set(directory, []);
333
+ }
334
+ filesByDirectory.get(directory)!.push(relativePath);
335
+ }
336
+ for (const [directory, directoryFiles] of filesByDirectory) {
337
+ if (directory === ".") {
338
+ for (const relativePath of directoryFiles) {
339
+ if (outputLines.length > 0) {
340
+ outputLines.push("");
341
+ }
342
+ outputLines.push(`# ${path.basename(relativePath)}`);
343
+ renderChunkedMatchesForFile(relativePath);
344
+ }
345
+ continue;
346
+ }
347
+ if (outputLines.length > 0) {
348
+ outputLines.push("");
349
+ }
350
+ outputLines.push(`# ${directory}`);
351
+ for (const relativePath of directoryFiles) {
352
+ outputLines.push(`## └─ ${path.basename(relativePath)}`);
353
+ renderChunkedMatchesForFile(relativePath);
354
+ }
355
+ }
356
+ } else {
357
+ for (const relativePath of fileList) {
358
+ renderChunkedMatchesForFile(relativePath);
359
+ }
360
+ }
361
+ const rawOutput = outputLines.join("\n");
293
362
  const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER });
294
363
  const truncated = Boolean(matchLimitReached || result.limitReached || truncation.truncated);
295
364
  const details: GrepToolDetails = {
@@ -1,6 +1,5 @@
1
1
  import type { AgentTool } from "@oh-my-pi/pi-agent-core";
2
2
  import type { ToolChoice } from "@oh-my-pi/pi-ai";
3
- import type { SearchDb } from "@oh-my-pi/pi-natives";
4
3
  import { $env, $flag, isBunTestRuntime, logger } from "@oh-my-pi/pi-utils";
5
4
  import type { AsyncJobManager } from "../async";
6
5
  import type { PromptTemplate } from "../config/prompt-templates";
@@ -8,7 +7,7 @@ import type { Settings } from "../config/settings";
8
7
  import { EditTool } from "../edit";
9
8
  import type { Skill } from "../extensibility/skills";
10
9
  import type { InternalUrlRouter } from "../internal-urls";
11
- import { getPreludeDocs, warmPythonEnvironment } from "../ipy/executor";
10
+ import { getPreludeDocs, resetPreludeDocsCache, warmPythonEnvironment } from "../ipy/executor";
12
11
  import { checkPythonKernelAvailability } from "../ipy/kernel";
13
12
  import { LspTool } from "../lsp";
14
13
  import type { DiscoverableMCPSearchIndex, DiscoverableMCPTool } from "../mcp/discoverable-tool-metadata";
@@ -22,10 +21,8 @@ import { SearchTool } from "../web/search";
22
21
  import { AskTool } from "./ask";
23
22
  import { AstEditTool } from "./ast-edit";
24
23
  import { AstGrepTool } from "./ast-grep";
25
- import { AwaitTool } from "./await-tool";
26
24
  import { BashTool } from "./bash";
27
25
  import { BrowserTool } from "./browser";
28
-
29
26
  import { CalculatorTool } from "./calculator";
30
27
  import { CancelJobTool } from "./cancel-job";
31
28
  import { type CheckpointState, CheckpointTool, RewindTool } from "./checkpoint";
@@ -47,6 +44,7 @@ import { GrepTool } from "./grep";
47
44
  import { InspectImageTool } from "./inspect-image";
48
45
  import { NotebookTool } from "./notebook";
49
46
  import { wrapToolWithMetaNotice } from "./output-meta";
47
+ import { PollTool } from "./poll-tool";
50
48
  import { PythonTool } from "./python";
51
49
  import { ReadTool } from "./read";
52
50
  import { RenderMermaidTool } from "./render-mermaid";
@@ -71,7 +69,6 @@ export * from "../web/search";
71
69
  export * from "./ask";
72
70
  export * from "./ast-edit";
73
71
  export * from "./ast-grep";
74
- export * from "./await-tool";
75
72
  export * from "./bash";
76
73
  export * from "./browser";
77
74
  export * from "./calculator";
@@ -85,6 +82,7 @@ export * from "./gh";
85
82
  export * from "./grep";
86
83
  export * from "./inspect-image";
87
84
  export * from "./notebook";
85
+ export * from "./poll-tool";
88
86
  export * from "./python";
89
87
  export * from "./read";
90
88
  export * from "./render-mermaid";
@@ -95,6 +93,7 @@ export * from "./search-tool-bm25";
95
93
  export * from "./ssh";
96
94
  export * from "./submit-result";
97
95
  export * from "./todo-write";
96
+ export * from "./vim";
98
97
  export * from "./write";
99
98
 
100
99
  /** Tool type (AgentTool from pi-ai) */
@@ -116,6 +115,8 @@ export interface ToolSession {
116
115
  hasUI: boolean;
117
116
  /** Skip Python kernel availability check and warmup */
118
117
  skipPythonPreflight?: boolean;
118
+ /** Force Python prelude warmup even when test env would normally skip it */
119
+ forcePythonWarmup?: boolean;
119
120
  /** Pre-loaded context files (AGENTS.md, etc) */
120
121
  contextFiles?: ContextFileEntry[];
121
122
  /** Pre-loaded skills */
@@ -124,7 +125,7 @@ export interface ToolSession {
124
125
  promptTemplates?: PromptTemplate[];
125
126
  /** Whether LSP integrations are enabled */
126
127
  enableLsp?: boolean;
127
- /** Whether the edit tool is available in this session (controls hashline output) */
128
+ /** Whether an edit-capable tool is available in this session (controls hashline output) */
128
129
  hasEditTool?: boolean;
129
130
  /** Event bus for tool/extension communication */
130
131
  eventBus?: EventBus;
@@ -136,6 +137,12 @@ export interface ToolSession {
136
137
  taskDepth?: number;
137
138
  /** Get session file */
138
139
  getSessionFile: () => string | null;
140
+ /** Get Python kernel owner ID for session-scoped retained-kernel cleanup */
141
+ getPythonKernelOwnerId?: () => string | null;
142
+ /** Reject new Python work once session disposal has started. */
143
+ assertPythonExecutionAllowed?: () => void;
144
+ /** Track tool-owned Python work so session disposal can await/abort it like direct session Python runs. */
145
+ trackPythonExecution?<T>(execution: Promise<T>, abortController: AbortController): Promise<T>;
139
146
  /** Get session ID */
140
147
  getSessionId?: () => string | null;
141
148
  /** Get artifacts directory for artifact:// URLs */
@@ -162,8 +169,6 @@ export interface ToolSession {
162
169
  asyncJobManager?: AsyncJobManager;
163
170
  /** Settings instance for passing to subagents */
164
171
  settings: Settings;
165
- /** Shared native search DB for grep/glob/fuzzyFind-backed workflows. */
166
- searchDb?: SearchDb;
167
172
  /** Plan mode state (if active) */
168
173
  getPlanModeState?: () => PlanModeState | undefined;
169
174
  /** Get compact conversation context for subagents (excludes tool results, system prompts) */
@@ -232,7 +237,7 @@ export const BUILTIN_TOOLS: Record<string, ToolFactory> = {
232
237
  rewind: RewindTool.createIf,
233
238
  task: TaskTool.create,
234
239
  cancel_job: CancelJobTool.createIf,
235
- await: AwaitTool.createIf,
240
+ poll: PollTool.createIf,
236
241
  todo_write: s => new TodoWriteTool(s),
237
242
  web_search: s => new SearchTool(s),
238
243
  search_tool_bm25: SearchToolBm25Tool.createIf,
@@ -297,7 +302,11 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P
297
302
  !skipPythonPreflight &&
298
303
  pythonMode !== "bash-only" &&
299
304
  (requestedTools === undefined || requestedTools.includes("python"));
300
- const skipPythonWarm = isBunTestRuntime() || $flag("PI_PYTHON_SKIP_CHECK");
305
+ const isTestEnv = isBunTestRuntime();
306
+ const forcePythonWarmup = session.forcePythonWarmup === true;
307
+ const skipPythonWarm = (isTestEnv && !forcePythonWarmup) || $flag("PI_PYTHON_SKIP_CHECK");
308
+ const cachedPreludeDocs = getPreludeDocs();
309
+ const shouldWarmPython = !skipPythonWarm && (forcePythonWarmup || cachedPreludeDocs.length === 0);
301
310
  if (shouldCheckPython) {
302
311
  const availability = await logger.time("createTools:pythonCheck", checkPythonKernelAvailability, session.cwd);
303
312
  pythonAvailable = availability.ok;
@@ -305,18 +314,38 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P
305
314
  logger.warn("Python kernel unavailable, falling back to bash", {
306
315
  reason: availability.reason,
307
316
  });
308
- } else if (!skipPythonWarm && getPreludeDocs().length === 0) {
317
+ } else if (shouldWarmPython) {
309
318
  const sessionFile = session.getSessionFile?.() ?? undefined;
319
+ const kernelOwnerId = session.getPythonKernelOwnerId?.() ?? undefined;
310
320
  const warmSessionId = sessionFile ? `session:${sessionFile}:cwd:${session.cwd}` : `cwd:${session.cwd}`;
321
+ const warmupAbortController = new AbortController();
311
322
  try {
312
- await logger.time(
313
- "createTools:warmPython",
314
- warmPythonEnvironment,
315
- session.cwd,
316
- warmSessionId,
317
- session.settings.get("python.sharedGateway"),
318
- sessionFile,
319
- );
323
+ session.assertPythonExecutionAllowed?.();
324
+ if (forcePythonWarmup && cachedPreludeDocs.length > 0) {
325
+ resetPreludeDocsCache();
326
+ }
327
+ const warmupExecution = session.trackPythonExecution
328
+ ? logger.time(
329
+ "createTools:warmPython",
330
+ warmPythonEnvironment,
331
+ session.cwd,
332
+ warmSessionId,
333
+ session.settings.get("python.sharedGateway"),
334
+ sessionFile,
335
+ kernelOwnerId,
336
+ warmupAbortController.signal,
337
+ )
338
+ : logger.time(
339
+ "createTools:warmPython",
340
+ warmPythonEnvironment,
341
+ session.cwd,
342
+ warmSessionId,
343
+ session.settings.get("python.sharedGateway"),
344
+ sessionFile,
345
+ kernelOwnerId,
346
+ );
347
+ await (session.trackPythonExecution?.(warmupExecution, warmupAbortController) ?? warmupExecution);
348
+ session.assertPythonExecutionAllowed?.();
320
349
  } catch (err) {
321
350
  logger.warn("Failed to warm Python environment", {
322
351
  error: err instanceof Error ? err.message : String(err),
@@ -2,10 +2,10 @@ import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallb
2
2
  import { prompt } from "@oh-my-pi/pi-utils";
3
3
  import { type Static, Type } from "@sinclair/typebox";
4
4
  import { isBackgroundJobSupportEnabled } from "../async";
5
- import awaitDescription from "../prompts/tools/await.md" with { type: "text" };
5
+ import pollDescription from "../prompts/tools/poll.md" with { type: "text" };
6
6
  import type { ToolSession } from "./index";
7
7
 
8
- const awaitSchema = Type.Object({
8
+ const pollSchema = Type.Object({
9
9
  jobs: Type.Optional(
10
10
  Type.Array(Type.String(), {
11
11
  description: "Specific job IDs to wait for. If omitted, waits for any running job.",
@@ -13,9 +13,9 @@ const awaitSchema = Type.Object({
13
13
  ),
14
14
  });
15
15
 
16
- type AwaitParams = Static<typeof awaitSchema>;
16
+ type PollParams = Static<typeof pollSchema>;
17
17
 
18
- interface AwaitResult {
18
+ interface PollResult {
19
19
  id: string;
20
20
  type: "bash" | "task";
21
21
  status: "running" | "completed" | "failed" | "cancelled";
@@ -25,33 +25,33 @@ interface AwaitResult {
25
25
  errorText?: string;
26
26
  }
27
27
 
28
- export interface AwaitToolDetails {
29
- jobs: AwaitResult[];
28
+ export interface PollToolDetails {
29
+ jobs: PollResult[];
30
30
  }
31
31
 
32
- export class AwaitTool implements AgentTool<typeof awaitSchema, AwaitToolDetails> {
33
- readonly name = "await";
34
- readonly label = "Await";
32
+ export class PollTool implements AgentTool<typeof pollSchema, PollToolDetails> {
33
+ readonly name = "poll";
34
+ readonly label = "Poll";
35
35
  readonly description: string;
36
- readonly parameters = awaitSchema;
36
+ readonly parameters = pollSchema;
37
37
  readonly strict = true;
38
38
 
39
39
  constructor(private readonly session: ToolSession) {
40
- this.description = prompt.render(awaitDescription);
40
+ this.description = prompt.render(pollDescription);
41
41
  }
42
42
 
43
- static createIf(session: ToolSession): AwaitTool | null {
43
+ static createIf(session: ToolSession): PollTool | null {
44
44
  if (!isBackgroundJobSupportEnabled(session.settings)) return null;
45
- return new AwaitTool(session);
45
+ return new PollTool(session);
46
46
  }
47
47
 
48
48
  async execute(
49
49
  _toolCallId: string,
50
- params: AwaitParams,
50
+ params: PollParams,
51
51
  signal?: AbortSignal,
52
- _onUpdate?: AgentToolUpdateCallback<AwaitToolDetails>,
52
+ _onUpdate?: AgentToolUpdateCallback<PollToolDetails>,
53
53
  _context?: AgentToolContext,
54
- ): Promise<AgentToolResult<AwaitToolDetails>> {
54
+ ): Promise<AgentToolResult<PollToolDetails>> {
55
55
  const manager = this.session.asyncJobManager;
56
56
  if (!manager) {
57
57
  return {
@@ -85,19 +85,25 @@ export class AwaitTool implements AgentTool<typeof awaitSchema, AwaitToolDetails
85
85
 
86
86
  // Block until at least one running job finishes or the call is aborted
87
87
  const racePromises: Promise<unknown>[] = runningJobs.map(j => j.promise);
88
-
89
- if (signal) {
90
- const { promise: abortPromise, resolve: abortResolve } = Promise.withResolvers<void>();
91
- const onAbort = () => abortResolve();
92
- signal.addEventListener("abort", onAbort, { once: true });
93
- racePromises.push(abortPromise);
94
- try {
88
+ const watchedJobIds = runningJobs.map(job => job.id);
89
+ manager.watchJobs(watchedJobIds);
90
+
91
+ try {
92
+ if (signal) {
93
+ const { promise: abortPromise, resolve: abortResolve } = Promise.withResolvers<void>();
94
+ const onAbort = () => abortResolve();
95
+ signal.addEventListener("abort", onAbort, { once: true });
96
+ racePromises.push(abortPromise);
97
+ try {
98
+ await Promise.race(racePromises);
99
+ } finally {
100
+ signal.removeEventListener("abort", onAbort);
101
+ }
102
+ } else {
95
103
  await Promise.race(racePromises);
96
- } finally {
97
- signal.removeEventListener("abort", onAbort);
98
104
  }
99
- } else {
100
- await Promise.race(racePromises);
105
+ } finally {
106
+ manager.unwatchJobs(watchedJobIds);
101
107
  }
102
108
 
103
109
  if (signal?.aborted) {
@@ -118,12 +124,12 @@ export class AwaitTool implements AgentTool<typeof awaitSchema, AwaitToolDetails
118
124
  resultText?: string;
119
125
  errorText?: string;
120
126
  }[],
121
- ): AgentToolResult<AwaitToolDetails> {
127
+ ): AgentToolResult<PollToolDetails> {
122
128
  const now = Date.now();
123
- const jobResults: AwaitResult[] = jobs.map(j => ({
129
+ const jobResults: PollResult[] = jobs.map(j => ({
124
130
  id: j.id,
125
131
  type: j.type,
126
- status: j.status as AwaitResult["status"],
132
+ status: j.status as PollResult["status"],
127
133
  label: j.label,
128
134
  durationMs: Math.max(0, now - j.startTime),
129
135
  ...(j.resultText ? { resultText: j.resultText } : {}),