pi-cursor-agent 0.1.0

package/src/pi/request-builder.ts

@@ -0,0 +1,279 @@
+import type {
+  Api,
+  Context,
+  Message,
+  Model,
+  TextContent,
+  Tool,
+  ToolResultMessage,
+} from "@mariozechner/pi-ai";
+import { type JsonValue, Value } from "@bufbuild/protobuf";
+import {
+  AgentClientMessage,
+  AgentConversationTurnStructure,
+  AgentRunRequest,
+  AssistantMessage as AssistantMessageProto,
+  ConversationAction,
+  type ConversationStateStructure,
+  ConversationStep,
+  ConversationTurnStructure,
+  ModelDetails,
+  UserMessage,
+  UserMessageAction,
+  ConversationStateStructure as ConversationStateStructureClass,
+} from "../__generated__/agent/v1/agent_pb";
+import {
+  type McpToolDefinition,
+  McpToolDefinition as McpToolDefinitionClass,
+  McpTools,
+} from "../__generated__/agent/v1/mcp_pb";
+import { toolResultToText } from "./utils/tool-result";
+import { getBlobId, type BlobStore } from "../vendor/agent-kv";
+
+const CURSOR_NATIVE_TOOL_NAMES = new Set([
+  "bash",
+  "read",
+  "write",
+  "delete",
+  "ls",
+  "grep",
+  "lsp",
+  "todo_write",
+]);
+
+type ContextWithTools = Context & { tools?: Tool[] };
+
+export function extractUserMessageText(msg: Message): string {
+  if (msg.role !== "user") return "";
+  if (typeof msg.content === "string") return msg.content.trim();
+  return msg.content
+    .filter((c): c is TextContent => c.type === "text")
+    .map((c) => c.text)
+    .join("\n")
+    .trim();
+}
+
+export function extractAssistantMessageText(msg: Message): string {
+  if (msg.role !== "assistant") return "";
+  if (!Array.isArray(msg.content)) return "";
+  return msg.content
+    .filter((c): c is TextContent => c.type === "text")
+    .map((c) => c.text)
+    .join("\n");
+}
+
+export function buildConversationTurns(messages: Message[]): Uint8Array[] {
+  const turns: Uint8Array[] = [];
+
+  let i = 0;
+  while (i < messages.length) {
+    const msg = messages[i];
+    if (!msg || msg.role !== "user") {
+      i++;
+      continue;
+    }
+
+    let isLastUserMessage = true;
+    for (let j = i + 1; j < messages.length; j++) {
+      if (messages[j]?.role === "user") {
+        isLastUserMessage = false;
+        break;
+      }
+    }
+    if (isLastUserMessage) break;
+
+    const userText = extractUserMessageText(msg);
+    if (!userText) {
+      i++;
+      continue;
+    }
+
+    const userMessage = new UserMessage({
+      text: userText,
+      messageId: crypto.randomUUID(),
+    });
+    const userMessageBytes = userMessage.toBinary();
+
+    const stepBytes: Uint8Array[] = [];
+    i++;
+    while (i < messages.length && messages[i]?.role !== "user") {
+      const stepMsg = messages[i];
+      if (!stepMsg) {
+        i++;
+        continue;
+      }
+
+      if (stepMsg.role === "assistant") {
+        const text = extractAssistantMessageText(stepMsg);
+        if (text) {
+          const step = new ConversationStep({
+            message: {
+              case: "assistantMessage",
+              value: new AssistantMessageProto({ text }),
+            },
+          });
+          stepBytes.push(step.toBinary());
+        }
+      } else if (stepMsg.role === "toolResult") {
+        const text = toolResultToText(stepMsg as ToolResultMessage);
+        if (text) {
+          const step = new ConversationStep({
+            message: {
+              case: "assistantMessage",
+              value: new AssistantMessageProto({
+                text: `[Tool Result]\n${text}`,
+              }),
+            },
+          });
+          stepBytes.push(step.toBinary());
+        }
+      }
+
+      i++;
+    }
+
+    const agentTurn = new AgentConversationTurnStructure({
+      userMessage: new Uint8Array(userMessageBytes),
+      steps: stepBytes,
+    });
+    const turn = new ConversationTurnStructure({
+      turn: { case: "agentConversationTurn", value: agentTurn },
+    });
+    turns.push(turn.toBinary());
+  }
+
+  return turns;
+}
+
+export function buildMcpToolDefinitions(
+  tools: Tool[] | undefined,
+): McpToolDefinition[] {
+  if (!tools || tools.length === 0) {
+    return [];
+  }
+
+  const advertisedTools = tools.filter(
+    (tool) => !CURSOR_NATIVE_TOOL_NAMES.has(tool.name),
+  );
+  if (advertisedTools.length === 0) {
+    return [];
+  }
+
+  return advertisedTools.map((tool) => {
+    const jsonSchema = tool.parameters as Record<string, unknown> | undefined;
+    const schemaValue: JsonValue =
+      jsonSchema && typeof jsonSchema === "object"
+        ? (jsonSchema as JsonValue)
+        : { type: "object", properties: {}, required: [] };
+    const inputSchema = new Uint8Array(Value.fromJson(schemaValue).toBinary());
+    return new McpToolDefinitionClass({
+      name: tool.name,
+      description: tool.description,
+      providerIdentifier: "pi-agent",
+      toolName: tool.name,
+      inputSchema,
+    });
+  });
+}
+
+export interface BuildRunRequestParams {
+  model: Model<Api>;
+  context: Context;
+  conversationId: string;
+  blobStore: BlobStore;
+  conversationState: ConversationStateStructure | undefined;
+  mcpToolDefinitions?: McpToolDefinition[];
+}
+
+export interface BuildRunRequestResult {
+  initialRequest: AgentClientMessage;
+  conversationState: ConversationStateStructure;
+}
+
+export function buildRunRequest(
+  params: BuildRunRequestParams,
+): BuildRunRequestResult {
+  const systemPromptJson = JSON.stringify({
+    role: "system",
+    content: params.context.systemPrompt || "You are a helpful assistant.",
+  });
+  const systemPromptBytes = new TextEncoder().encode(systemPromptJson);
+  const systemPromptId = getBlobId(systemPromptBytes);
+  void params.blobStore.setBlob(null, systemPromptId, systemPromptBytes);
+
+  const lastMessage = params.context.messages.at(-1);
+  const userText = lastMessage ? extractUserMessageText(lastMessage) : "";
+  if (!userText) {
+    throw new Error("Cannot send empty user message to Cursor API");
+  }
+
+  const userMessage = new UserMessage({
+    text: userText,
+    messageId: crypto.randomUUID(),
+  });
+
+  const action = new ConversationAction({
+    action: {
+      case: "userMessageAction",
+      value: new UserMessageAction({ userMessage }),
+    },
+  });
+
+  const cached = params.conversationState;
+  const hasMatchingPrompt = cached?.rootPromptMessagesJson?.some((entry) =>
+    Buffer.from(entry).equals(systemPromptId),
+  );
+
+  const turns = buildConversationTurns(params.context.messages);
+
+  const baseState =
+    cached && hasMatchingPrompt
+      ? cached
+      : new ConversationStateStructureClass({
+          rootPromptMessagesJson: [systemPromptId],
+          turns: [],
+          todos: [],
+          pendingToolCalls: [],
+          previousWorkspaceUris: [],
+          fileStates: {},
+          fileStatesV2: {},
+          summaryArchives: [],
+          turnTimings: [],
+          subagentStates: {},
+          selfSummaryCount: 0,
+          readPaths: [],
+        });
+
+  const conversationState = new ConversationStateStructureClass({
+    ...baseState,
+    turns: turns.length > 0 ? turns : baseState.turns,
+  });
+
+  const modelDetails = new ModelDetails({
+    modelId: params.model.id,
+    displayModelId: params.model.id,
+    displayName: params.model.name,
+  });
+
+  const mcpToolDefinitions = params.mcpToolDefinitions ?? [];
+  const runRequest = new AgentRunRequest({
+    conversationState,
+    action,
+    modelDetails,
+    conversationId: params.conversationId,
+    mcpTools: new McpTools({ mcpTools: mcpToolDefinitions }),
+  });
+
+  const initialRequest = new AgentClientMessage({
+    message: { case: "runRequest", value: runRequest },
+  });
+
+  return {
+    initialRequest,
+    conversationState,
+  };
+}
+
+export function getContextTools(context: Context): McpToolDefinition[] {
+  return buildMcpToolDefinitions((context as ContextWithTools).tools);
+}

package/src/pi/utils/tool-result.ts

@@ -0,0 +1,35 @@
+import type {
+  ImageContent,
+  TextContent,
+  ToolResultMessage,
+} from "@mariozechner/pi-ai";
+
+export function toolResultToText(result: ToolResultMessage): string {
+  return result.content
+    .map((item: TextContent | ImageContent) => {
+      if (item.type === "text") return item.text;
+      return `[${item.mimeType} image]`;
+    })
+    .join("\n");
+}
+
+export function toolResultWasTruncated(result: ToolResultMessage): boolean {
+  if (!result.details || typeof result.details !== "object") {
+    return false;
+  }
+  const truncation = (
+    result.details as { truncation?: { truncated?: boolean } }
+  ).truncation;
+  return !!truncation?.truncated;
+}
+
+export function toolResultDetailBoolean(
+  result: ToolResultMessage,
+  key: string,
+): boolean {
+  if (!result.details || typeof result.details !== "object") {
+    return false;
+  }
+  const value = (result.details as Record<string, unknown>)[key];
+  return typeof value === "boolean" ? value : false;
+}

package/src/stream.ts

@@ -0,0 +1,386 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
+import {
+  createAssistantMessageEventStream,
+  type Api,
+  type AssistantMessage,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+  type TextContent,
+  type ThinkingContent,
+} from "@mariozechner/pi-ai";
+import type { ConversationStateStructure } from "./__generated__/agent/v1/agent_pb";
+import {
+  AskQuestionRejected,
+  AskQuestionResult,
+} from "./__generated__/agent/v1/ask_question_tool_pb";
+import AgentService from "./api/agent-service";
+import { toCursorId } from "./pi/model-mapping";
+import { buildRunRequest, getContextTools } from "./pi/request-builder";
+import {
+  CURSOR_STATE_ENTRY_TYPE,
+  ensureAgentStore,
+  persistAgentStore,
+} from "./pi/agent-store";
+import {
+  LocalResourceProvider,
+  type PiToolContext,
+  type ToolExecEvent,
+} from "./pi/local-resource-provider";
+import {
+  AgentConnectClient,
+  type CheckpointHandler,
+  type InteractionListener,
+} from "./vendor/agent-client";
+import type {
+  CoreInteractionUpdate,
+  CoreInteractionQuery,
+  CoreInteractionResponse,
+} from "./vendor/agent-core";
+
+export const CURSOR_API_URL = "https://api2.cursor.sh";
+const CURSOR_CLIENT_VERSION = "cli-2026.01.17-d239e66";
+
+function createCheckpointHandler(
+  handler: (checkpoint: ConversationStateStructure) => void,
+): CheckpointHandler {
+  return {
+    handleCheckpoint(
+      _ctx: unknown,
+      checkpoint: ConversationStateStructure,
+    ): Promise<void> {
+      handler(checkpoint);
+      return Promise.resolve();
+    },
+  };
+}
+
+const QUERY_REJECTION_REASON = "Not supported by pi-cursor-agent";
+
+function createInteractionListenerAdapter(
+  onUpdate: (update: CoreInteractionUpdate) => void,
+): InteractionListener {
+  return {
+    async sendUpdate(
+      _ctx: unknown,
+      update: CoreInteractionUpdate,
+    ): Promise<void> {
+      onUpdate(update);
+    },
+
+    async query(
+      _ctx: unknown,
+      query: CoreInteractionQuery,
+    ): Promise<CoreInteractionResponse> {
+      switch (query.type) {
+        case "ask-question-request":
+          return {
+            result: new AskQuestionResult({
+              result: {
+                case: "rejected",
+                value: new AskQuestionRejected({
+                  reason: QUERY_REJECTION_REASON,
+                }),
+              },
+            }),
+          };
+        case "web-search-request":
+        case "web-fetch-request":
+        case "exa-search-request":
+        case "exa-fetch-request":
+        case "switch-mode-request":
+          return { approved: false, reason: QUERY_REJECTION_REASON };
+        case "create-plan-request":
+          return {
+            result: {
+              planUri: "",
+              result: {
+                case: "error",
+                value: { error: QUERY_REJECTION_REASON },
+              },
+            },
+          } as CoreInteractionResponse;
+        case "setup-vm-environment-request":
+          return {} as CoreInteractionResponse;
+        default:
+          return { approved: false, reason: QUERY_REJECTION_REASON };
+      }
+    },
+  };
+}
+
+export function streamCursorAgent(
+  pi: ExtensionAPI,
+  getCtx: () => ExtensionContext | null,
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+): AssistantMessageEventStream {
+  const stream = createAssistantMessageEventStream();
+
+  (async () => {
+    const startTime = Date.now();
+    let firstTokenTime: number | undefined;
+
+    const output: AssistantMessage = {
+      role: "assistant",
+      content: [],
+      api: model.api,
+      provider: model.provider,
+      model: model.id,
+      usage: {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+      },
+      stopReason: "stop",
+      timestamp: Date.now(),
+    };
+
+    const sessionId = options?.sessionId ?? "default";
+
+    try {
+      const apiKey = options?.apiKey;
+      if (!apiKey) {
+        throw new Error(
+          "Cursor API key (access token) is required. Run /login cursor or set CURSOR_ACCESS_TOKEN.",
+        );
+      }
+
+      const agentStore = await ensureAgentStore(sessionId);
+      const cwd = getCtx()?.cwd ?? process.cwd();
+      const requestContextTools = getContextTools(context);
+
+      let onToolExec: ((event: ToolExecEvent) => void) | undefined;
+
+      const piToolCtx: PiToolContext = {
+        cwd,
+        ...(options?.signal ? { signal: options.signal } : {}),
+        getActiveTools: () => new Set(pi.getActiveTools()),
+        getCtx,
+        onToolExec: (event) => onToolExec?.(event),
+      };
+
+      const resources = new LocalResourceProvider({
+        ctx: piToolCtx,
+        requestContextTools,
+      });
+
+      const blobStore = agentStore.getBlobStore();
+      const cursorModelId = toCursorId(model.id, options?.reasoning);
+      const { initialRequest, conversationState } = buildRunRequest({
+        model: { ...model, id: cursorModelId },
+        context,
+        conversationId: agentStore.getId(),
+        blobStore,
+        conversationState: agentStore.getConversationStateStructure(),
+        mcpToolDefinitions: requestContextTools,
+      });
+      agentStore.conversationStateStructure = conversationState;
+
+      stream.push({ type: "start", partial: output });
+
+      let currentTextBlock: TextContent | null = null;
+      let currentThinkingBlock: ThinkingContent | null = null;
+      const usageState = { sawTokenDelta: false };
+
+      const finalizeTextBlock = () => {
+        if (!currentTextBlock) return;
+        stream.push({
+          type: "text_end",
+          contentIndex: output.content.indexOf(currentTextBlock),
+          content: currentTextBlock.text,
+          partial: output,
+        });
+        currentTextBlock = null;
+      };
+
+      const finalizeThinkingBlock = () => {
+        if (!currentThinkingBlock) return;
+        stream.push({
+          type: "thinking_end",
+          contentIndex: output.content.indexOf(currentThinkingBlock),
+          content: currentThinkingBlock.thinking,
+          partial: output,
+        });
+        currentThinkingBlock = null;
+      };
+
+      onToolExec = (event: ToolExecEvent) => {
+        if (event.type === "start") {
+          finalizeTextBlock();
+          finalizeThinkingBlock();
+          (stream as any).push({
+            type: "tool_exec_start",
+            toolCallId: event.toolCallId,
+            toolName: event.toolName,
+            args: event.args,
+          });
+        } else {
+          (stream as any).push({
+            type: "tool_exec_end",
+            toolCallId: event.toolCallId,
+            toolName: event.toolName,
+            result: {
+              content: event.result.content,
+              details: event.result.details,
+            },
+            isError: event.result.isError,
+          });
+        }
+      };
+
+      const handleInteractionUpdate = (update: CoreInteractionUpdate) => {
+        switch (update.type) {
+          case "text-delta": {
+            if (!firstTokenTime) firstTokenTime = Date.now();
+            finalizeThinkingBlock();
+            const delta = update.text;
+            if (!currentTextBlock) {
+              currentTextBlock = { type: "text", text: "" };
+              output.content.push(currentTextBlock);
+              stream.push({
+                type: "text_start",
+                contentIndex: output.content.length - 1,
+                partial: output,
+              });
+            }
+            currentTextBlock.text += delta;
+            stream.push({
+              type: "text_delta",
+              contentIndex: output.content.indexOf(currentTextBlock),
+              delta,
+              partial: output,
+            });
+            return;
+          }
+
+          case "thinking-delta": {
+            if (!firstTokenTime) firstTokenTime = Date.now();
+            finalizeTextBlock();
+            const delta = update.text;
+            if (!currentThinkingBlock) {
+              currentThinkingBlock = { type: "thinking", thinking: "" };
+              output.content.push(currentThinkingBlock);
+              stream.push({
+                type: "thinking_start",
+                contentIndex: output.content.length - 1,
+                partial: output,
+              });
+            }
+            currentThinkingBlock.thinking += delta;
+            stream.push({
+              type: "thinking_delta",
+              contentIndex: output.content.indexOf(currentThinkingBlock),
+              delta,
+              partial: output,
+            });
+            return;
+          }
+
+          case "thinking-completed": {
+            finalizeThinkingBlock();
+            return;
+          }
+
+          case "turn-ended": {
+            output.stopReason = "stop";
+            return;
+          }
+
+          case "token-delta": {
+            usageState.sawTokenDelta = true;
+            output.usage.output += update.tokens;
+            output.usage.totalTokens = output.usage.input + output.usage.output;
+            return;
+          }
+        }
+      };
+
+      const baseUrl = model.baseUrl || CURSOR_API_URL;
+      const agentService = new AgentService(baseUrl, {
+        accessToken: apiKey,
+        clientVersion: CURSOR_CLIENT_VERSION,
+        clientType: "cli",
+      });
+
+      const connectClient = new AgentConnectClient(agentService.rpcClient);
+
+      const interactionListener = createInteractionListenerAdapter(
+        handleInteractionUpdate,
+      );
+
+      const checkpointHandler = createCheckpointHandler(
+        (checkpoint: ConversationStateStructure) => {
+          void agentStore.handleCheckpoint(null, checkpoint);
+          if (usageState.sawTokenDelta) return;
+          const usedTokens = checkpoint.tokenDetails?.usedTokens ?? 0;
+          if (usedTokens > 0 && output.usage.output !== usedTokens) {
+            output.usage.output = usedTokens;
+            output.usage.totalTokens = output.usage.input + output.usage.output;
+          }
+        },
+      );
+      checkpointHandler.getLatestCheckpoint = () =>
+        agentStore.getConversationStateStructure();
+
+      const runOptions: Parameters<typeof connectClient.run>[1] = {
+        interactionListener,
+        resources,
+        blobStore,
+        checkpointHandler,
+      };
+      if (options?.signal) runOptions.signal = options.signal;
+
+      await connectClient.run(initialRequest, runOptions);
+
+      finalizeTextBlock();
+      finalizeThinkingBlock();
+
+      output.usage.cost = {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        total: 0,
+      };
+      (output as any).duration = Date.now() - startTime;
+      if (firstTokenTime) {
+        (output as any).ttft = firstTokenTime - startTime;
+      }
+
+      stream.push({ type: "done", reason: "stop", message: output });
+      stream.end();
+    } catch (error) {
+      output.stopReason = options?.signal?.aborted ? "aborted" : "error";
+      output.errorMessage =
+        error instanceof Error ? error.message : String(error);
+      (output as any).duration = Date.now() - startTime;
+      if (firstTokenTime) (output as any).ttft = firstTokenTime - startTime;
+      stream.push({
+        type: "error",
+        reason: output.stopReason as any,
+        error: output,
+      });
+      stream.end();
+    } finally {
+      try {
+        const snapshot = await persistAgentStore(sessionId);
+        if (snapshot) {
+          pi.appendEntry(CURSOR_STATE_ENTRY_TYPE, snapshot);
+        }
+      } catch {
+        // ignore persistence errors
+      }
+    }
+  })();
+
+  return stream;
+}