@posthog/agent
Version:
TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog
656 lines (639 loc) • 27.8 kB
JavaScript
import { randomBytes } from 'node:crypto';
import { McpServer } from '../../../node_modules/@modelcontextprotocol/sdk/dist/esm/server/mcp.js';
import { createPatch } from '../../../node_modules/diff/libesm/patch/create.js';
import { z } from 'zod';
import { Logger } from '../../utils/logger.js';
import { extractLinesWithByteLimit, sleep, unreachable } from './utils.js';
const SYSTEM_REMINDER = `
<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>`;
const defaults = { maxFileSize: 50000, linesToRead: 2000 };
const unqualifiedToolNames = {
read: "Read",
edit: "Edit",
write: "Write",
bash: "Bash",
killShell: "KillShell",
bashOutput: "BashOutput",
};
const SERVER_PREFIX = "mcp__acp__";
const toolNames = {
read: SERVER_PREFIX + unqualifiedToolNames.read,
edit: SERVER_PREFIX + unqualifiedToolNames.edit,
write: SERVER_PREFIX + unqualifiedToolNames.write,
bash: SERVER_PREFIX + unqualifiedToolNames.bash,
killShell: SERVER_PREFIX + unqualifiedToolNames.killShell,
bashOutput: SERVER_PREFIX + unqualifiedToolNames.bashOutput,
};
const EDIT_TOOL_NAMES = [toolNames.edit, toolNames.write];
function createMcpServer(agent, sessionId, clientCapabilities) {
// Create MCP server
const server = new McpServer({ name: "acp", version: "1.0.0" }, { capabilities: { tools: {} } });
if (clientCapabilities?.fs?.readTextFile) {
server.registerTool(unqualifiedToolNames.read, {
title: unqualifiedToolNames.read,
description: `Reads the content of the given file in the project.
In sessions with ${toolNames.read} always use it instead of Read as it contains the most up-to-date contents.
Reads a file from the local filesystem. If the User provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned.
Usage:
- The file_path parameter must be an absolute path, not a relative path
- By default, it reads up to ${defaults.linesToRead} lines starting from the beginning of the file
- You can optionally specify a line offset and limit (especially handy for long files), but it's recommended to read the whole file by not providing these parameters
- Any files larger than ${defaults.maxFileSize} bytes will be truncated
- This tool allows Claude Code to read images (eg PNG, JPG, etc). When reading an image file the contents are presented visually as Claude Code is a multimodal LLM.
- This tool can only read files, not directories. To read a directory, use an ls command via the ${toolNames.bash} tool.
- You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful.`,
inputSchema: {
file_path: z
.string()
.describe("The absolute path to the file to read"),
offset: z
.number()
.optional()
.default(1)
.describe("The line number to start reading from. Only provide if the file is too large to read at once"),
limit: z
.number()
.optional()
.default(defaults.linesToRead)
.describe(`The number of lines to read. Only provide if the file is too large to read at once.`),
},
annotations: {
title: "Read file",
readOnlyHint: true,
destructiveHint: false,
openWorldHint: false,
idempotentHint: false,
},
}, async (input) => {
try {
const session = agent.sessions[sessionId];
if (!session) {
return {
content: [
{
type: "text",
text: "The user has left the building",
},
],
};
}
const readResponse = await agent.readTextFile({
sessionId,
path: input.file_path,
line: input.offset,
limit: input.limit,
});
if (typeof readResponse?.content !== "string") {
throw new Error(`No file contents for ${input.file_path}.`);
}
// Extract lines with byte limit enforcement
const result = extractLinesWithByteLimit(readResponse.content, defaults.maxFileSize);
// Construct informative message about what was read
let readInfo = "";
if (input.offset > 1 || result.wasLimited) {
readInfo = "\n\n<file-read-info>";
if (result.wasLimited) {
readInfo += `Read ${result.linesRead} lines (hit 50KB limit). `;
}
else {
readInfo += `Read lines ${input.offset}-${result.linesRead}. `;
}
if (result.wasLimited) {
readInfo += `Continue with offset=${result.linesRead}.`;
}
readInfo += "</file-read-info>";
}
return {
content: [
{
type: "text",
text: result.content + readInfo + SYSTEM_REMINDER,
},
],
};
}
catch (error) {
return {
content: [
{
type: "text",
text: `Reading file failed: ${error.message}`,
},
],
};
}
});
}
if (clientCapabilities?.fs?.writeTextFile) {
server.registerTool(unqualifiedToolNames.write, {
title: unqualifiedToolNames.write,
description: `Writes a file to the local filesystem..
In sessions with ${toolNames.write} always use it instead of Write as it will
allow the user to conveniently review changes.
Usage:
- This tool will overwrite the existing file if there is one at the provided path.
- If this is an existing file, you MUST use the ${toolNames.read} tool first to read the file's contents. This tool will fail if you did not read the file first.
- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required.
- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.
- Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.`,
inputSchema: {
file_path: z
.string()
.describe("The absolute path to the file to write (must be absolute, not relative)"),
content: z.string().describe("The content to write to the file"),
},
annotations: {
title: "Write file",
readOnlyHint: false,
destructiveHint: false,
openWorldHint: false,
idempotentHint: false,
},
}, async (input) => {
try {
const session = agent.sessions[sessionId];
if (!session) {
return {
content: [
{
type: "text",
text: "The user has left the building",
},
],
};
}
await agent.writeTextFile({
sessionId,
path: input.file_path,
content: input.content,
});
return {
content: [],
};
}
catch (error) {
return {
content: [
{
type: "text",
text: `Writing file failed: ${error.message}`,
},
],
};
}
});
server.registerTool(unqualifiedToolNames.edit, {
title: unqualifiedToolNames.edit,
description: `Performs exact string replacements in files.
In sessions with ${toolNames.edit} always use it instead of Edit as it will
allow the user to conveniently review changes.
Usage:
- You must use your \`${toolNames.read}\` tool at least once in the conversation before editing. This tool will error if you attempt an edit without reading the file.
- When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears.
- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required.
- Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked.
- The edit will FAIL if \`old_string\` is not unique in the file. Either provide a larger string with more surrounding context to make it unique or use \`replace_all\` to change every instance of \`old_string\`.
- Use \`replace_all\` for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.`,
inputSchema: {
file_path: z
.string()
.describe("The absolute path to the file to modify"),
old_string: z.string().describe("The text to replace"),
new_string: z
.string()
.describe("The text to replace it with (must be different from old_string)"),
replace_all: z
.boolean()
.default(false)
.optional()
.describe("Replace all occurences of old_string (default false)"),
},
annotations: {
title: "Edit file",
readOnlyHint: false,
destructiveHint: false,
openWorldHint: false,
idempotentHint: false,
},
}, async (input) => {
try {
const session = agent.sessions[sessionId];
if (!session) {
return {
content: [
{
type: "text",
text: "The user has left the building",
},
],
};
}
const readResponse = await agent.readTextFile({
sessionId,
path: input.file_path,
});
if (typeof readResponse?.content !== "string") {
throw new Error(`No file contents for ${input.file_path}.`);
}
const { newContent } = replaceAndCalculateLocation(readResponse.content, [
{
oldText: input.old_string,
newText: input.new_string,
replaceAll: input.replace_all,
},
]);
const patch = createPatch(input.file_path, readResponse.content, newContent);
await agent.writeTextFile({
sessionId,
path: input.file_path,
content: newContent,
});
return {
content: [
{
type: "text",
text: patch,
},
],
};
}
catch (error) {
return {
content: [
{
type: "text",
text: `Editing file failed: ${error?.message ?? String(error)}`,
},
],
};
}
});
}
if (agent.clientCapabilities?.terminal) {
server.registerTool(unqualifiedToolNames.bash, {
title: unqualifiedToolNames.bash,
description: `Executes a bash command
In sessions with ${toolNames.bash} always use it instead of Bash`,
inputSchema: {
command: z.string().describe("The command to execute"),
timeout: z
.number()
.default(2 * 60 * 1000)
.describe(`Optional timeout in milliseconds (max ${2 * 60 * 1000})`),
description: z
.string()
.optional()
.describe(`Clear, concise description of what this command does in 5-10 words, in active voice. Examples:
Input: ls
Output: List files in current directory
Input: git status
Output: Show working tree status
Input: npm install
Output: Install package dependencies
Input: mkdir foo
Output: Create directory 'foo'`),
run_in_background: z
.boolean()
.default(false)
.describe(`Set to true to run this command in the background. The tool returns an \`id\` that can be used with the \`${toolNames.bashOutput}\` tool to retrieve the current output, or the \`${toolNames.killShell}\` tool to stop it early.`),
},
}, async (input, extra) => {
const session = agent.sessions[sessionId];
if (!session) {
return {
content: [
{
type: "text",
text: "The user has left the building",
},
],
};
}
const toolCallId = extra._meta?.["claudecode/toolUseId"];
if (typeof toolCallId !== "string") {
throw new Error("No tool call ID found");
}
if (!agent.clientCapabilities?.terminal ||
!agent.client.createTerminal) {
throw new Error("unreachable");
}
const handle = await agent.client.createTerminal({
command: input.command,
env: [{ name: "CLAUDECODE", value: "1" }],
sessionId,
outputByteLimit: 32_000,
});
await agent.client.sessionUpdate({
sessionId,
update: {
sessionUpdate: "tool_call_update",
toolCallId,
status: "in_progress",
title: input.description,
content: [{ type: "terminal", terminalId: handle.id }],
},
});
const abortPromise = new Promise((resolve) => {
if (extra.signal.aborted) {
resolve(null);
}
else {
extra.signal.addEventListener("abort", () => {
resolve(null);
});
}
});
const statusPromise = Promise.race([
handle
.waitForExit()
.then((exitStatus) => ({ status: "exited", exitStatus })),
abortPromise.then(() => ({
status: "aborted",
exitStatus: null,
})),
sleep(input.timeout).then(async () => {
if (agent.backgroundTerminals[handle.id]?.status === "started") {
await handle.kill();
}
return { status: "timedOut", exitStatus: null };
}),
]);
if (input.run_in_background) {
agent.backgroundTerminals[handle.id] = {
handle,
lastOutput: null,
status: "started",
};
statusPromise.then(async ({ status, exitStatus }) => {
const bgTerm = agent.backgroundTerminals[handle.id];
if (bgTerm.status !== "started") {
return;
}
const currentOutput = await handle.currentOutput();
agent.backgroundTerminals[handle.id] = {
status,
pendingOutput: {
...currentOutput,
output: stripCommonPrefix(bgTerm.lastOutput?.output ?? "", currentOutput.output),
exitStatus: exitStatus ?? currentOutput.exitStatus,
},
};
return handle.release();
});
return {
content: [
{
type: "text",
text: `Command started in background with id: ${handle.id}`,
},
],
};
}
await using terminal = handle;
const { status } = await statusPromise;
if (status === "aborted") {
return {
content: [{ type: "text", text: "Tool cancelled by user" }],
};
}
const output = await terminal.currentOutput();
return {
content: [{ type: "text", text: toolCommandOutput(status, output) }],
};
});
server.registerTool(unqualifiedToolNames.bashOutput, {
title: unqualifiedToolNames.bashOutput,
description: `- Retrieves output from a running or completed background bash shell
- Takes a shell_id parameter identifying the shell
- Always returns only new output since the last check
- Returns stdout and stderr output along with shell status
- Use this tool when you need to monitor or check the output of a long-running shell
In sessions with ${toolNames.bashOutput} always use it instead of BashOutput.`,
inputSchema: {
shell_id: z
.string()
.describe(`The id of the background bash command as returned by \`${toolNames.bash}\``),
},
}, async (input) => {
const bgTerm = agent.backgroundTerminals[input.shell_id];
if (!bgTerm) {
throw new Error(`Unknown shell ${input.shell_id}`);
}
if (bgTerm.status === "started") {
const newOutput = await bgTerm.handle.currentOutput();
const strippedOutput = stripCommonPrefix(bgTerm.lastOutput?.output ?? "", newOutput.output);
bgTerm.lastOutput = newOutput;
return {
content: [
{
type: "text",
text: toolCommandOutput(bgTerm.status, {
...newOutput,
output: strippedOutput,
}),
},
],
};
}
else {
return {
content: [
{
type: "text",
text: toolCommandOutput(bgTerm.status, bgTerm.pendingOutput),
},
],
};
}
});
server.registerTool(unqualifiedToolNames.killShell, {
title: unqualifiedToolNames.killShell,
description: `- Kills a running background bash shell by its ID
- Takes a shell_id parameter identifying the shell to kill
- Returns a success or failure status
- Use this tool when you need to terminate a long-running shell
In sessions with ${toolNames.killShell} always use it instead of KillShell.`,
inputSchema: {
shell_id: z
.string()
.describe(`The id of the background bash command as returned by \`${toolNames.bash}\``),
},
}, async (input) => {
const bgTerm = agent.backgroundTerminals[input.shell_id];
if (!bgTerm) {
throw new Error(`Unknown shell ${input.shell_id}`);
}
switch (bgTerm.status) {
case "started": {
await bgTerm.handle.kill();
const currentOutput = await bgTerm.handle.currentOutput();
agent.backgroundTerminals[bgTerm.handle.id] = {
status: "killed",
pendingOutput: {
...currentOutput,
output: stripCommonPrefix(bgTerm.lastOutput?.output ?? "", currentOutput.output),
},
};
await bgTerm.handle.release();
return {
content: [{ type: "text", text: "Command killed successfully." }],
};
}
case "aborted":
return {
content: [{ type: "text", text: "Command aborted by user." }],
};
case "exited":
return {
content: [{ type: "text", text: "Command had already exited." }],
};
case "killed":
return {
content: [{ type: "text", text: "Command was already killed." }],
};
case "timedOut":
return {
content: [{ type: "text", text: "Command killed by timeout." }],
};
default: {
unreachable(bgTerm, new Logger({ prefix: "[McpServer]" }));
throw new Error("Unexpected background terminal status");
}
}
});
}
return server;
}
function stripCommonPrefix(a, b) {
let i = 0;
while (i < a.length && i < b.length && a[i] === b[i]) {
i++;
}
return b.slice(i);
}
function toolCommandOutput(status, output) {
const { exitStatus, output: commandOutput, truncated } = output;
let toolOutput = "";
switch (status) {
case "started":
case "exited": {
if (exitStatus && (exitStatus.exitCode ?? null) === null) {
toolOutput += `Interrupted by the user. `;
}
break;
}
case "killed":
toolOutput += `Killed. `;
break;
case "timedOut":
toolOutput += `Timed out. `;
break;
case "aborted":
break;
default: {
const unreachable = status;
return unreachable;
}
}
if (exitStatus) {
if (typeof exitStatus.exitCode === "number") {
toolOutput += `Exited with code ${exitStatus.exitCode}.`;
}
if (typeof exitStatus.signal === "string") {
toolOutput += `Signal \`${exitStatus.signal}\`. `;
}
toolOutput += "Final output:\n\n";
}
else {
toolOutput += "New output:\n\n";
}
toolOutput += commandOutput;
if (truncated) {
toolOutput += `\n\nCommand output was too long, so it was truncated to ${commandOutput.length} bytes.`;
}
return toolOutput;
}
/**
* Replace text in a file and calculate the line numbers where the edits occurred.
*
* @param fileContent - The full file content
* @param edits - Array of edit operations to apply sequentially
* @returns the new content and the line numbers where replacements occurred in the final content
*/
function replaceAndCalculateLocation(fileContent, edits) {
let currentContent = fileContent;
// Use unique markers to track where replacements happen
const markerPrefix = `__REPLACE_MARKER_${randomBytes(5).toString("hex")}_`;
let markerCounter = 0;
const markers = [];
// Apply edits sequentially, inserting markers at replacement positions
for (const edit of edits) {
// Skip empty oldText
if (edit.oldText === "") {
throw new Error(`The provided \`old_string\` is empty.\n\nNo edits were applied.`);
}
if (edit.replaceAll) {
// Replace all occurrences with marker + newText
const parts = [];
let lastIndex = 0;
let searchIndex = 0;
while (true) {
const index = currentContent.indexOf(edit.oldText, searchIndex);
if (index === -1) {
if (searchIndex === 0) {
throw new Error(`The provided \`old_string\` does not appear in the file: "${edit.oldText}".\n\nNo edits were applied.`);
}
break;
}
// Add content before the match
parts.push(currentContent.substring(lastIndex, index));
// Add marker and replacement
const marker = `${markerPrefix}${markerCounter++}__`;
markers.push(marker);
parts.push(marker + edit.newText);
lastIndex = index + edit.oldText.length;
searchIndex = lastIndex;
}
// Add remaining content
parts.push(currentContent.substring(lastIndex));
currentContent = parts.join("");
}
else {
// Replace first occurrence only
const index = currentContent.indexOf(edit.oldText);
if (index === -1) {
throw new Error(`The provided \`old_string\` does not appear in the file: "${edit.oldText}".\n\nNo edits were applied.`);
}
else {
const marker = `${markerPrefix}${markerCounter++}__`;
markers.push(marker);
currentContent =
currentContent.substring(0, index) +
marker +
edit.newText +
currentContent.substring(index + edit.oldText.length);
}
}
}
// Find line numbers where markers appear in the content
const lineNumbers = [];
for (const marker of markers) {
const index = currentContent.indexOf(marker);
if (index !== -1) {
const lineNumber = Math.max(0, currentContent.substring(0, index).split(/\r\n|\r|\n/).length - 1);
lineNumbers.push(lineNumber);
}
}
// Remove all markers from the final content
let finalContent = currentContent;
for (const marker of markers) {
finalContent = finalContent.replace(marker, "");
}
// Dedupe and sort line numbers
const uniqueLineNumbers = [...new Set(lineNumbers)].sort();
return { newContent: finalContent, lineNumbers: uniqueLineNumbers };
}
export { EDIT_TOOL_NAMES, SYSTEM_REMINDER, createMcpServer, replaceAndCalculateLocation, toolNames };
//# sourceMappingURL=mcp-server.js.map