106 lines
3.1 KiB
TypeScript
106 lines
3.1 KiB
TypeScript
/**
|
|
* Atuin extension for pi.
|
|
*
|
|
* Tracks bash commands executed by pi in Atuin history with author `pi`.
|
|
*
|
|
* Install with:
|
|
* atuin hook install pi
|
|
*
|
|
* Then restart pi or run /reload.
|
|
*/
|
|
|
|
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
|
|
|
|
const ATUIN_AUTHOR = "pi";
|
|
const ATUIN_TIMEOUT_MS = 10_000;
|
|
|
|
async function startHistory(
|
|
pi: ExtensionAPI,
|
|
cwd: string,
|
|
command: string,
|
|
): Promise<string | undefined> {
|
|
try {
|
|
const result = await pi.exec(
|
|
"atuin",
|
|
["history", "start", "--author", ATUIN_AUTHOR, "--", command],
|
|
{ cwd, timeout: ATUIN_TIMEOUT_MS },
|
|
);
|
|
|
|
if (result.code !== 0) return undefined;
|
|
|
|
const id = result.stdout.trim();
|
|
return id.length > 0 ? id : undefined;
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
async function endHistory(
|
|
pi: ExtensionAPI,
|
|
cwd: string,
|
|
historyId: string,
|
|
exitCode: number,
|
|
): Promise<void> {
|
|
try {
|
|
await pi.exec(
|
|
"atuin",
|
|
["history", "end", historyId, "--exit", String(exitCode)],
|
|
{ cwd, timeout: ATUIN_TIMEOUT_MS },
|
|
);
|
|
} catch {
|
|
// Ignore Atuin failures so command execution is never blocked.
|
|
}
|
|
}
|
|
|
|
// The bash tool reports failures by appending a status line to the result
|
|
// text rather than exposing a numeric exit code, so recover it from there.
|
|
function exitCodeFromResult(result: unknown, isError: boolean): number {
|
|
if (!isError) return 0;
|
|
|
|
const content = (result as { content?: unknown } | undefined)?.content;
|
|
const text = Array.isArray(content)
|
|
? content
|
|
.map((part) => {
|
|
const t = (part as { text?: unknown } | undefined)?.text;
|
|
return typeof t === "string" ? t : "";
|
|
})
|
|
.join("\n")
|
|
: "";
|
|
|
|
const exited = text.match(/Command exited with code (\d+)\s*$/);
|
|
if (exited) return Number(exited[1]);
|
|
if (/Command aborted\s*$/.test(text)) return 130;
|
|
if (/Command timed out after \S+ seconds\s*$/.test(text)) return 124;
|
|
return 1;
|
|
}
|
|
|
|
export default function atuinPiExtension(pi: ExtensionAPI) {
|
|
// Atuin history IDs for in-flight bash tool calls, keyed by tool call ID.
|
|
const pending = new Map<string, string>();
|
|
|
|
// Observe bash executions through events instead of registering a bash
|
|
// tool: registering one conflicts with other extensions that provide
|
|
// their own bash tool (sandboxes, RTK, remote runners), while events
|
|
// fire no matter which extension's bash tool ends up executing the
|
|
// command.
|
|
pi.on("tool_call", async (event, ctx: ExtensionContext) => {
|
|
if (event.toolName !== "bash") return;
|
|
|
|
const command = (event.input as { command?: unknown }).command;
|
|
if (typeof command !== "string" || command.length === 0) return;
|
|
|
|
const historyId = await startHistory(pi, ctx.cwd, command);
|
|
if (historyId) pending.set(event.toolCallId, historyId);
|
|
});
|
|
|
|
// tool_execution_end also fires when another extension blocks the call,
|
|
// unlike tool_result, so entries started above are always closed.
|
|
pi.on("tool_execution_end", async (event, ctx: ExtensionContext) => {
|
|
const historyId = pending.get(event.toolCallId);
|
|
if (!historyId) return;
|
|
pending.delete(event.toolCallId);
|
|
|
|
await endHistory(pi, ctx.cwd, historyId, exitCodeFromResult(event.result, event.isError));
|
|
});
|
|
}
|