pi-cursor-agent 0.1.0

package/src/api/ai-service.ts

@@ -0,0 +1,42 @@
+import {
+  createClient,
+  type Client,
+  type Interceptor,
+} from "@connectrpc/connect";
+import { createConnectTransport } from "@connectrpc/connect-node";
+import { AiService as AiServiceDef } from "../__generated__/aiserver/v1/aiserver_service_connect";
+
+export interface AgentServiceOptions {
+  accessToken: string;
+  clientType: string;
+  clientVersion: string;
+}
+
+class AiService {
+  private readonly client: Client<typeof AiServiceDef>;
+
+  constructor(baseUrl: string, options: AgentServiceOptions) {
+    const authInterceptor: Interceptor = (next) => async (req) => {
+      req.header.set("authorization", `Bearer ${options.accessToken}`);
+      req.header.set("x-cursor-client-type", options.clientType);
+      req.header.set("x-cursor-client-version", options.clientVersion);
+      req.header.set("x-ghost-mode", "true");
+      req.header.set("x-request-id", crypto.randomUUID());
+      return next(req);
+    };
+
+    const transport = createConnectTransport({
+      baseUrl,
+      httpVersion: "1.1",
+      interceptors: [authInterceptor],
+    });
+
+    this.client = createClient(AiServiceDef, transport);
+  }
+
+  public async getUsableModels() {
+    return this.client.getUsableModels({});
+  }
+}
+
+export default AiService;

package/src/api/auth.ts

@@ -0,0 +1,74 @@
+type AuthResult = {
+  accessToken: string;
+  refreshToken: string;
+};
+
+class Auth {
+  constructor(private readonly baseUrl: string) {}
+
+  public async poll({
+    uuid,
+    verifier,
+    signal,
+  }: {
+    uuid: string;
+    verifier: string;
+    signal?: AbortSignal | undefined;
+  }) {
+    const params = new URLSearchParams({ uuid, verifier });
+    return this.fetchJson<AuthResult>(`/auth/poll?${params.toString()}`, {
+      headers: { "content-type": "application/json" },
+      signal: signal ?? null,
+      validator: this.isAuthResult,
+    });
+  }
+
+  public async exchangeUserApiKey({
+    token,
+    signal,
+  }: {
+    token: string;
+    signal?: AbortSignal | undefined;
+  }) {
+    return this.fetchJson<AuthResult>("/auth/exchange_user_api_key", {
+      method: "POST",
+      headers: {
+        authorization: `Bearer ${token}`,
+        "content-type": "application/json",
+      },
+      body: JSON.stringify({}),
+      signal: signal ?? null,
+      validator: this.isAuthResult,
+    });
+  }
+
+  private async fetchJson<T>(
+    url: string,
+    { validator, ...init }: RequestInit & { validator: (data: T) => data is T },
+  ): Promise<T> {
+    const response = await fetch(`${this.baseUrl}${url}`, init);
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`Fetch failed ${url} for ${response.status}: ${error}`);
+    }
+
+    const data = await response.json();
+    if (!validator(data)) {
+      const error = JSON.stringify(data);
+      throw new Error(`Fetch failed ${url} for invalid response: ${error}`);
+    }
+
+    return data;
+  }
+
+  private isAuthResult(data: unknown): data is AuthResult {
+    return (
+      typeof data === "object" &&
+      data !== null &&
+      "accessToken" in data &&
+      "refreshToken" in data
+    );
+  }
+}
+
+export default Auth;

package/src/index.ts

@@ -0,0 +1,101 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
+import type {
+  Api,
+  OAuthCredentials,
+  OAuthLoginCallbacks,
+} from "@mariozechner/pi-ai";
+import AiService from "./api/ai-service";
+import Auth from "./api/auth";
+import AuthManager from "./lib/auth";
+import { restoreAgentStoreFromBranch } from "./pi/agent-store";
+import { streamCursorAgent } from "./stream";
+import {
+  CURSOR_API_URL,
+  CURSOR_CLIENT_VERSION,
+  CURSOR_WEBSITE_URL,
+} from "./lib/env";
+import { getCachedPiModels, updateCachedPiModels } from "./pi/model";
+
+const auth = new AuthManager(new Auth(CURSOR_API_URL), CURSOR_WEBSITE_URL);
+
+const login = async (
+  callbacks: OAuthLoginCallbacks,
+): Promise<OAuthCredentials> => {
+  const credentials = await auth.login(callbacks);
+  const ai = new AiService(CURSOR_API_URL, {
+    accessToken: credentials.access,
+    clientVersion: CURSOR_CLIENT_VERSION,
+    clientType: "cli",
+  });
+  await updateCachedPiModels(ai);
+  return credentials;
+};
+
+const refreshToken = async (
+  credentials: OAuthCredentials,
+): Promise<OAuthCredentials> => {
+  const refreshed = await auth.refresh(credentials);
+  const ai = new AiService(CURSOR_API_URL, {
+    accessToken: credentials.access,
+    clientVersion: CURSOR_CLIENT_VERSION,
+    clientType: "cli",
+  });
+  await updateCachedPiModels(ai);
+  return refreshed;
+};
+
+export default (pi: ExtensionAPI) => {
+  let lastCtx: ExtensionContext | null = null;
+  const getCtx = () => lastCtx;
+
+  const refreshBranchState = async (ctx: ExtensionContext) => {
+    lastCtx = ctx;
+    try {
+      const sessionId = ctx.sessionManager.getSessionId();
+      await restoreAgentStoreFromBranch(
+        sessionId,
+        ctx.sessionManager.getBranch(),
+      );
+    } catch {
+      // ignore
+    }
+  };
+
+  pi.on("session_start", async (_, ctx) => {
+    await refreshBranchState(ctx);
+  });
+
+  pi.on("session_switch", async (_, ctx) => {
+    await refreshBranchState(ctx);
+  });
+
+  pi.on("session_tree", async (_, ctx) => {
+    await refreshBranchState(ctx);
+  });
+
+  pi.on("before_agent_start", async (_, ctx) => {
+    lastCtx = ctx;
+  });
+
+  pi.on("agent_start", async (_, ctx) => {
+    lastCtx = ctx;
+  });
+
+  pi.registerProvider("cursor", {
+    baseUrl: CURSOR_API_URL,
+    apiKey: "CURSOR_ACCESS_TOKEN",
+    api: "cursor-agent" as unknown as Api,
+    streamSimple: (model, context, options) =>
+      streamCursorAgent(pi, getCtx, model, context, options),
+    models: getCachedPiModels(),
+    oauth: {
+      name: "Cursor Agent",
+      login,
+      refreshToken,
+      getApiKey: (cred) => cred.access,
+    },
+  });
+};

package/src/lib/agent-store/disk.ts

@@ -0,0 +1,139 @@
+import fs from "node:fs/promises";
+import path from "node:path";
+import { toHex, fromHex, type AgentMetadata } from "../../vendor/agent-kv";
+
+interface BlobEntry {
+  id: string;
+  data: string;
+}
+
+interface BlobsFile {
+  version: 1;
+  blobs: BlobEntry[];
+}
+
+interface MetaFile {
+  version: 1;
+  agentId: string;
+  latestRootBlobId: string;
+  name: string;
+  createdAt: number;
+  mode: string;
+  lastUsedModel?: string;
+}
+
+const getSessionDir = (baseDir: string, sessionId: string): string =>
+  path.join(baseDir, "chats", sessionId);
+
+const getBlobsFilePath = (baseDir: string, sessionId: string): string =>
+  path.join(getSessionDir(baseDir, sessionId), "blobs.json");
+
+const getMetaFilePath = (baseDir: string, sessionId: string): string =>
+  path.join(getSessionDir(baseDir, sessionId), "meta.json");
+
+export const loadBlobsFromDisk = async (
+  baseDir: string,
+  sessionId: string,
+): Promise<Map<string, Uint8Array>> => {
+  try {
+    const text = await fs.readFile(
+      getBlobsFilePath(baseDir, sessionId),
+      "utf-8",
+    );
+    const parsed = JSON.parse(text) as BlobsFile;
+    if (!parsed || parsed.version !== 1 || !Array.isArray(parsed.blobs)) {
+      return new Map();
+    }
+    const map = new Map<string, Uint8Array>();
+    for (const entry of parsed.blobs) {
+      if (
+        !entry ||
+        typeof entry.id !== "string" ||
+        typeof entry.data !== "string"
+      )
+        continue;
+      try {
+        map.set(entry.id, new Uint8Array(Buffer.from(entry.data, "base64")));
+      } catch {
+        // skip corrupt entries
+      }
+    }
+    return map;
+  } catch {
+    return new Map();
+  }
+};
+
+export const saveBlobsToDisk = async (
+  baseDir: string,
+  sessionId: string,
+  blobs: Map<string, Uint8Array>,
+): Promise<void> => {
+  const dir = getSessionDir(baseDir, sessionId);
+  await fs.mkdir(dir, { recursive: true });
+  const file: BlobsFile = {
+    version: 1,
+    blobs: Array.from(blobs.entries()).map(([id, data]) => ({
+      id,
+      data: Buffer.from(data).toString("base64"),
+    })),
+  };
+  const filePath = getBlobsFilePath(baseDir, sessionId);
+  const tmpPath = `${filePath}.tmp`;
+  await fs.writeFile(tmpPath, JSON.stringify(file), "utf-8");
+  await fs.rename(tmpPath, filePath);
+};
+
+export const loadMetaFromDisk = async (
+  baseDir: string,
+  sessionId: string,
+): Promise<AgentMetadata | null> => {
+  try {
+    const text = await fs.readFile(
+      getMetaFilePath(baseDir, sessionId),
+      "utf-8",
+    );
+    const parsed = JSON.parse(text) as MetaFile;
+    if (!parsed || parsed.version !== 1 || typeof parsed.agentId !== "string") {
+      return null;
+    }
+    return {
+      agentId: parsed.agentId,
+      latestRootBlobId: parsed.latestRootBlobId
+        ? fromHex(parsed.latestRootBlobId)
+        : new Uint8Array(),
+      name: parsed.name ?? "New Agent",
+      createdAt: parsed.createdAt ?? Date.now(),
+      mode: (parsed.mode as AgentMetadata["mode"]) ?? "default",
+      ...(parsed.lastUsedModel != null && {
+        lastUsedModel: parsed.lastUsedModel,
+      }),
+    };
+  } catch {
+    return null;
+  }
+};
+
+export const saveMetaToDisk = async (
+  baseDir: string,
+  sessionId: string,
+  metadata: AgentMetadata,
+): Promise<void> => {
+  const dir = getSessionDir(baseDir, sessionId);
+  await fs.mkdir(dir, { recursive: true });
+  const file: MetaFile = {
+    version: 1,
+    agentId: metadata.agentId,
+    latestRootBlobId: toHex(metadata.latestRootBlobId),
+    name: metadata.name,
+    createdAt: metadata.createdAt,
+    mode: metadata.mode,
+    ...(metadata.lastUsedModel != null && {
+      lastUsedModel: metadata.lastUsedModel,
+    }),
+  };
+  const filePath = getMetaFilePath(baseDir, sessionId);
+  const tmpPath = `${filePath}.tmp`;
+  await fs.writeFile(tmpPath, JSON.stringify(file), "utf-8");
+  await fs.rename(tmpPath, filePath);
+};

package/src/lib/agent-store/index.ts

@@ -0,0 +1,72 @@
+import { AgentStore, getDefaultAgentMetadata } from "../../vendor/agent-kv";
+import {
+  loadBlobsFromDisk,
+  loadMetaFromDisk,
+  saveBlobsToDisk,
+  saveMetaToDisk,
+} from "./disk";
+import { JsonBlobStoreWithMetadata } from "./json-blob-store";
+
+export interface StoreEntry {
+  store: AgentStore;
+  jsonStore: JsonBlobStoreWithMetadata;
+}
+
+const sessionStores = new Map<string, StoreEntry>();
+
+export const ensureAgentStore = async (
+  baseDir: string,
+  sessionId: string,
+): Promise<StoreEntry> => {
+  const existing = sessionStores.get(sessionId);
+  if (existing) {
+    return existing;
+  }
+
+  const [blobs, meta] = await Promise.all([
+    loadBlobsFromDisk(baseDir, sessionId),
+    loadMetaFromDisk(baseDir, sessionId),
+  ]);
+
+  const metadata = meta ?? getDefaultAgentMetadata();
+  const jsonStore = new JsonBlobStoreWithMetadata(blobs, metadata);
+  const store = new AgentStore(jsonStore, jsonStore);
+
+  if (metadata.latestRootBlobId.length > 0) {
+    await store.resetFromDb(null);
+  }
+
+  const entry: StoreEntry = { store, jsonStore };
+  sessionStores.set(sessionId, entry);
+  return entry;
+};
+
+export const persistAgentStore = async (
+  baseDir: string,
+  sessionId: string,
+): Promise<StoreEntry | null> => {
+  const entry = sessionStores.get(sessionId);
+  if (!entry) {
+    return null;
+  }
+
+  await Promise.all([
+    saveBlobsToDisk(baseDir, sessionId, entry.jsonStore.blobs),
+    saveMetaToDisk(baseDir, sessionId, entry.jsonStore.metadata),
+  ]);
+
+  return entry;
+};
+
+export const applySnapshotToStore = async (
+  entry: StoreEntry,
+  agentId: string,
+  latestRootBlobId: Uint8Array,
+): Promise<void> => {
+  entry.jsonStore.metadata.agentId = agentId;
+  entry.jsonStore.metadata.latestRootBlobId = latestRootBlobId;
+
+  if (latestRootBlobId.length > 0) {
+    await entry.store.resetFromDb(null);
+  }
+};

package/src/lib/agent-store/json-blob-store.ts

@@ -0,0 +1,47 @@
+import {
+  toHex,
+  type AgentMetadata,
+  type BlobStore,
+  type MetadataStore,
+} from "../../vendor/agent-kv";
+
+export class JsonBlobStoreWithMetadata implements BlobStore, MetadataStore {
+  readonly blobs: Map<string, Uint8Array>;
+  readonly metadata: AgentMetadata;
+
+  constructor(blobs: Map<string, Uint8Array>, metadata: AgentMetadata) {
+    this.blobs = blobs;
+    this.metadata = metadata;
+  }
+
+  public get<K extends keyof AgentMetadata>(key: K): AgentMetadata[K] {
+    return this.metadata[key];
+  }
+
+  public set<K extends keyof AgentMetadata>(
+    key: K,
+    value: AgentMetadata[K],
+  ): void {
+    this.metadata[key] = value;
+  }
+
+  public subscribe(_: keyof AgentMetadata, __: () => void): () => void {
+    return () => {};
+  }
+
+  public async getBlob(
+    _: unknown,
+    blobId: Uint8Array,
+  ): Promise<Uint8Array | undefined> {
+    return Promise.resolve(this.blobs.get(toHex(blobId)));
+  }
+
+  public setBlob(
+    _ctx: unknown,
+    blobId: Uint8Array,
+    blobData: Uint8Array,
+  ): Promise<void> {
+    this.blobs.set(toHex(blobId), blobData);
+    return Promise.resolve();
+  }
+}

package/src/lib/auth.ts

@@ -0,0 +1,135 @@
+import { createHash, randomBytes } from "node:crypto";
+import { OAuthLoginCallbacks } from "@mariozechner/pi-ai";
+import Auth from "../api/auth";
+import { backoff } from "./backoff";
+
+type OnAuth = (info: { url: string; instructions: string }) => void;
+type OnProgress = (message: string) => void;
+
+class AuthManager {
+  constructor(
+    private readonly auth: Auth,
+    private readonly websiteUrl: string,
+  ) {}
+
+  public async login({
+    onAuth,
+    onProgress,
+    signal,
+  }: {
+    onAuth: OnAuth;
+    onProgress?: OnProgress;
+    signal?: AbortSignal;
+  }) {
+    const { uuid, verifier, loginUrl } = this.generateAuthParams();
+
+    const instructions = "Complete the sign-in in your browser.";
+    onAuth({ url: loginUrl, instructions });
+
+    return await this.pollAuthenticationStatus({
+      uuid,
+      verifier,
+      onProgress,
+      signal,
+    });
+  }
+
+  public async refresh(credentials: {
+    access: string;
+    refresh: string;
+  }): Promise<{ access: string; refresh: string; expires: number }> {
+    if (!credentials.access && !credentials.refresh) {
+      throw new Error("No credentials provided");
+    }
+
+    try {
+      const { accessToken, refreshToken } = await this.auth.exchangeUserApiKey({
+        token: credentials.refresh || credentials.access,
+      });
+      const expires = getTokenExpiry(accessToken);
+      return { access: accessToken, refresh: refreshToken, expires };
+    } catch {
+      // If the refresh token is invalid, try to refresh it with access token
+      if (credentials.access && credentials.refresh) {
+        return this.refresh({ access: credentials.access, refresh: "" });
+      }
+      throw new Error("Failed to refresh credentials");
+    }
+  }
+
+  private generateAuthParams() {
+    const verifier = base64URLEncode(randomBytes(32));
+    const challenge = base64URLEncode(
+      createHash("sha256").update(verifier).digest(),
+    );
+    const uuid = crypto.randomUUID();
+    const loginUrl = `${this.websiteUrl}/loginDeepControl?challenge=${challenge}&uuid=${uuid}&mode=login&redirectTarget=cli`;
+    return { challenge, uuid, verifier, loginUrl };
+  }
+
+  private async pollAuthenticationStatus({
+    uuid,
+    verifier,
+    onProgress,
+    signal,
+  }: {
+    uuid: string;
+    verifier: string;
+    onProgress?: OAuthLoginCallbacks["onProgress"];
+    signal?: AbortSignal | undefined;
+  }) {
+    return backoff(
+      async () => {
+        onProgress?.("Polling authentication status...");
+        const tokens = await this.auth.poll({ uuid, verifier, signal });
+        const { accessToken, refreshToken } = tokens;
+        const expires = getTokenExpiry(accessToken);
+        return { access: accessToken, refresh: refreshToken, expires };
+      },
+      {
+        retries: 150,
+        delay: 1000,
+        shouldRetry: (error) =>
+          error instanceof Error &&
+          error.message.includes("/auth/poll") &&
+          error.message.includes("404"),
+      },
+    );
+  }
+}
+
+type JwtPayload = {
+  exp: number;
+  [key: string]: unknown;
+};
+
+export const base64URLEncode = (buffer: Buffer) => {
+  return buffer
+    .toString("base64")
+    .replace(/\+/g, "-")
+    .replace(/\//g, "_")
+    .replace(/=/g, "");
+};
+
+export const decodeJwt = (token: string): JwtPayload | null => {
+  try {
+    const parts = token.split(".");
+    if (parts.length !== 3) return null;
+    const payload = parts[1] ?? "";
+    const decoded = atob(payload);
+    return JSON.parse(decoded) as JwtPayload;
+  } catch {
+    return null;
+  }
+};
+
+const getTokenExpiry = (token: string): number => {
+  try {
+    const decoded = decodeJwt(token);
+    return decoded!.exp * 1000 - 5 * 60 * 1000;
+  } catch {
+    return Date.now() + 3600 * 1000;
+  }
+};
+
+export default AuthManager;

package/src/lib/backoff.ts

@@ -0,0 +1,32 @@
+type BackoffConfig = {
+  retries: number;
+  delay: number;
+  shouldRetry?: (error: unknown, attempt: number) => boolean;
+};
+
+const sleep = (ms: number): Promise<void> => {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+};
+
+export const backoff = async <T>(
+  fn: () => Promise<T>,
+  { retries, delay, shouldRetry }: BackoffConfig,
+): Promise<T> => {
+  let lastError: unknown;
+
+  for (let attempt = 0; attempt < retries; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error;
+
+      if (shouldRetry && !shouldRetry(error, attempt)) {
+        throw error;
+      }
+
+      await sleep(delay);
+    }
+  }
+
+  throw lastError;
+};

package/src/lib/env.ts

@@ -0,0 +1,3 @@
+export const CURSOR_API_URL = "https://api2.cursor.sh";
+export const CURSOR_WEBSITE_URL = "https://cursor.com";
+export const CURSOR_CLIENT_VERSION = "cli-2026.01.17-d239e66";

package/src/lib/heartbeat.ts

@@ -0,0 +1,21 @@
+export const heartbeat = (
+  fn: () => Promise<unknown>,
+  interval: number,
+): (() => void) => {
+  let timeout: ReturnType<typeof setTimeout> | undefined;
+
+  const schedule = () => {
+    timeout = setTimeout(() => {
+      fn().then(schedule, () => {});
+    }, interval);
+  };
+
+  schedule();
+
+  return () => {
+    if (timeout !== undefined) {
+      clearTimeout(timeout);
+      timeout = undefined;
+    }
+  };
+};