char-slop/ai-dots

ai dotfiles

git clone https://git.t4t.associates/char-slop/ai-dots

Charlotte Sompi: compact reasoning, computer-usec535dce

main
6.8 KiB216 linesraw
1// raw thinking notepad (based on github:lyramakesmusic/neuralese-leaker)
2
3import { Type } from "@earendil-works/pi-ai";
4import { type ExtensionAPI, getMarkdownTheme } from "@earendil-works/pi-coding-agent";
5import { Container, Markdown } from "@earendil-works/pi-tui";
6import { loadSettings } from "./_char_common/settings";
7
8const TOOL_NAME = "notepad";
9const PARAMETER_IDLE_TIMEOUT_MS = 10_000;
10
11const SCRATCH_REPLY = "Continue thinking, call a tool, or respond to the user.";
12
13const SYSTEM_PROMPT_ADDENDUM = `# notepad
14
15- Call notepad first, before anything else, on every new task.
16- Re-enter the notepad after every tool result, before your next action.`;
17
18export default function (pi: ExtensionAPI) {
19  let enabled = true;
20  let watchedCallId: string | undefined;
21  let idleTimer: ReturnType<typeof setTimeout> | undefined;
22  let resumeTimer: ReturnType<typeof setTimeout> | undefined;
23  let resumeAfterAgentEnd = false;
24  const truncatedCalls = new Set<string>();
25
26  function clearIdleTimer() {
27    if (idleTimer) clearTimeout(idleTimer);
28    idleTimer = undefined;
29    watchedCallId = undefined;
30  }
31
32  function isActive() {
33    return enabled && pi.getActiveTools().includes(TOOL_NAME);
34  }
35
36  function setActive(on: boolean) {
37    enabled = on;
38    if (!on) clearIdleTimer();
39    const active = pi.getActiveTools();
40    if (on && !active.includes(TOOL_NAME)) {
41      pi.setActiveTools([...active, TOOL_NAME]);
42    } else if (!on && active.includes(TOOL_NAME)) {
43      pi.setActiveTools(active.filter((t) => t !== TOOL_NAME));
44    }
45  }
46
47  pi.registerTool({
48    name: TOOL_NAME,
49    label: "Notepad",
50    description:
51      "Notepad: Unlike typical scratchpads, it has no budget cap. " +
52      "Use it before and between actions.",
53    promptSnippet: "notepad; think here before and between actions",
54    promptGuidelines: [
55      "Use notepad to plan before acting, and re-enter it after every tool result.",
56    ],
57    parameters: Type.Object({
58      text: Type.String({
59        description: "Working notes",
60      }),
61    }),
62    renderShell: "self",
63    async execute(toolCallId) {
64      const terminate = truncatedCalls.delete(toolCallId);
65      return {
66        content: [{ type: "text", text: SCRATCH_REPLY }],
67        details: {},
68        ...(terminate ? { terminate: true } : {}),
69      };
70    },
71    renderCall(args, theme, context) {
72      const text = typeof args?.text === "string" ? args.text : "";
73      const md =
74        (context.lastComponent as Markdown | undefined) ??
75        new Markdown("", 1, 0, getMarkdownTheme(), {
76          color: (t) => theme.fg("thinkingText", t),
77          italic: true,
78        });
79      md.setText(text);
80      return md;
81    },
82    renderResult: () => new Container(),
83  });
84
85  pi.on("session_start", async (_event, ctx) => {
86    const tools = loadSettings(ctx.cwd).tools;
87    const fromSettings =
88      tools && typeof tools === "object"
89        ? (tools as Record<string, unknown>).notepad
90        : undefined;
91    enabled = fromSettings !== false;
92    if (enabled) setActive(true);
93  });
94
95  pi.on("message_update", (event, ctx) => {
96    if (!isActive()) return;
97    const update = event.assistantMessageEvent;
98    if (
99      update.type !== "toolcall_start" &&
100      update.type !== "toolcall_delta" &&
101      update.type !== "toolcall_end"
102    ) {
103      return;
104    }
105
106    const call = update.partial.content[update.contentIndex];
107    if (call?.type !== "toolCall" || call.name !== TOOL_NAME) return;
108    if (update.type === "toolcall_end") {
109      clearIdleTimer();
110      return;
111    }
112
113    clearIdleTimer();
114    watchedCallId = call.id;
115    idleTimer = setTimeout(() => {
116      if (watchedCallId !== call.id) return;
117      clearIdleTimer();
118      truncatedCalls.add(call.id);
119      resumeAfterAgentEnd = true;
120      ctx.abort();
121    }, PARAMETER_IDLE_TIMEOUT_MS);
122  });
123
124  // An aborted response normally skips tool execution. Reclassify only our
125  // timed-out call so its partial arguments become an ordinary tool result.
126  pi.on("message_end", (event) => {
127    clearIdleTimer();
128    if (event.message.role !== "assistant" || event.message.stopReason !== "aborted") return;
129
130    const truncated = event.message.content.some(
131      (part) => part.type === "toolCall" && truncatedCalls.has(part.id),
132    );
133    if (!truncated) return;
134
135    const { errorMessage: _errorMessage, ...message } = event.message;
136    return {
137      message: {
138        ...message,
139        stopReason: "toolUse" as const,
140        content: message.content.map((part) =>
141          part.type === "toolCall" && truncatedCalls.has(part.id)
142            ? {
143                ...part,
144                arguments: {
145                  ...part.arguments,
146                  work: typeof part.arguments.work === "string" ? part.arguments.work : "",
147                },
148              }
149            : part,
150        ),
151      },
152    };
153  });
154
155  // The salvaged result terminates the aborted run; resume once its signal is
156  // gone so the next provider request gets a fresh one.
157  pi.on("agent_end", (_event, ctx) => {
158    if (!resumeAfterAgentEnd) return;
159    resumeAfterAgentEnd = false;
160
161    const resume = () => {
162      if (!ctx.isIdle()) {
163        resumeTimer = setTimeout(resume, 10);
164        return;
165      }
166      resumeTimer = undefined;
167      pi.sendMessage(
168        {
169          customType: "notepad-timeout",
170          content:
171            "The previous notepad call was truncated after its arguments stopped streaming. Continue, calling notepad again if more reasoning is needed.",
172          display: false,
173        },
174        { triggerTurn: true },
175      );
176    };
177    resumeTimer = setTimeout(resume, 0);
178  });
179
180  pi.on("session_shutdown", () => {
181    clearIdleTimer();
182    if (resumeTimer) clearTimeout(resumeTimer);
183    resumeTimer = undefined;
184    resumeAfterAgentEnd = false;
185    truncatedCalls.clear();
186  });
187
188  pi.on("before_agent_start", async (event) => {
189    if (!isActive()) return;
190    pi.setThinkingLevel("off");
191    if (event.systemPrompt.includes("# raw thinking (notepad)")) return;
192    return { systemPrompt: event.systemPrompt + "\n\n" + SYSTEM_PROMPT_ADDENDUM };
193  });
194
195  pi.registerCommand("notepad", {
196    description: "Toggle raw-thinking notepad: /notepad [on|off|status]",
197    getArgumentCompletions: (prefix) => {
198      const normalized = (prefix || "").trim().toLowerCase();
199      const items = ["on", "off", "status"].filter((v) => v.startsWith(normalized));
200      return items.length > 0 ? items.map((value) => ({ value, label: value })) : null;
201    },
202    handler: async (args, ctx) => {
203      const arg = (args || "").trim().toLowerCase();
204      if (arg === "" || arg === "status") {
205        ctx.ui.notify(`Notepad raw thinking: ${isActive() ? "on" : "off"}`, "info");
206        return;
207      }
208      if (arg !== "on" && arg !== "off") {
209        ctx.ui.notify("Usage: /notepad [on|off|status]", "warning");
210        return;
211      }
212      setActive(arg === "on");
213      ctx.ui.notify(`Notepad raw thinking: ${arg} (this session)`, "info");
214    },
215  });
216}