@posthog/agent 2.3.366 → 2.3.385

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/agent",
-  "version": "2.3.366",
+  "version": "2.3.385",
   "repository": "https://github.com/PostHog/code",
   "description": "TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog",
   "exports": {

package/src/adapters/claude/claude-agent.ts

@@ -316,15 +316,6 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
   }
 
   async prompt(params: PromptRequest): Promise<PromptResponse> {
-    this.session.cancelled = false;
-    this.session.interruptReason = undefined;
-    this.session.accumulatedUsage = {
-      inputTokens: 0,
-      outputTokens: 0,
-      cachedReadTokens: 0,
-      cachedWriteTokens: 0,
-    };
-
     const userMessage = promptToClaude(params);
     const promptUuid = randomUUID();
     userMessage.uuid = promptUuid;
@@ -364,7 +355,19 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
       this.session.input.push(userMessage);
     }
 
-    // Broadcast user message to client
+    // Reset session state here (after the queued-wait) rather than at the
+    // top of prompt(). Otherwise a new prompt() call would wipe cancelled=true
+    // on the previous still-running loop, causing it to return end_turn
+    // instead of the cancelled stop reason the spec requires.
+    this.session.cancelled = false;
+    this.session.interruptReason = undefined;
+    this.session.accumulatedUsage = {
+      inputTokens: 0,
+      outputTokens: 0,
+      cachedReadTokens: 0,
+      cachedWriteTokens: 0,
+    };
+
     await this.broadcastUserMessage(params);
 
     this.session.promptRunning = true;
@@ -652,10 +655,6 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
 
           case "user":
           case "assistant": {
-            if (this.session.cancelled) {
-              break;
-            }
-
             // Check for prompt replay (our own message echoed back)
             if (message.type === "user" && "uuid" in message && message.uuid) {
               if (message.uuid === promptUuid) {
@@ -670,12 +669,16 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
                 pending.resolve(false);
                 this.session.pendingMessages.delete(message.uuid as string);
                 handedOff = true;
-                // the current loop stops with end_turn,
-                // the loop of the next prompt continues running
-                return { stopReason: "end_turn" };
+                return {
+                  stopReason: this.session.cancelled ? "cancelled" : "end_turn",
+                };
               }
             }
 
+            if (this.session.cancelled) {
+              break;
+            }
+
             // Skip replayed user messages that aren't pending prompts
             if (
               "isReplay" in message &&

package/src/enrichment/file-enricher.test.ts

@@ -8,11 +8,15 @@ function makeDeps(overrides: {
   parseRejects?: Error;
   isSupported?: boolean;
   getApiKey?: () => string | Promise<string>;
+  findImportsInSource?: () => Promise<unknown[]>;
+  getWrappersForFile?: () => Promise<unknown[]>;
 }): {
   deps: FileEnrichmentDeps;
   parseSpy: ReturnType<typeof vi.fn>;
   enrichFromApiSpy: ReturnType<typeof vi.fn>;
   getApiKeySpy: ReturnType<typeof vi.fn>;
+  findImportsSpy: ReturnType<typeof vi.fn>;
+  getWrappersSpy: ReturnType<typeof vi.fn>;
 } {
   const enrichFromApiSpy = vi.fn(async () => ({
     toInlineComments: () =>
@@ -29,11 +33,19 @@ function makeDeps(overrides: {
   });
 
   const getApiKeySpy = vi.fn(overrides.getApiKey ?? (() => "phx_test"));
+  const findImportsSpy = vi.fn(
+    overrides.findImportsInSource ?? (async () => []),
+  );
+  const getWrappersSpy = vi.fn(
+    overrides.getWrappersForFile ?? (async () => []),
+  );
 
   const deps: FileEnrichmentDeps = {
     enricher: {
       isSupported: vi.fn(() => overrides.isSupported ?? true),
       parse: parseSpy,
+      findImportsInSource: findImportsSpy,
+      getWrappersForFile: getWrappersSpy,
     } as unknown as FileEnrichmentDeps["enricher"],
     apiConfig: {
       apiUrl: "https://test.posthog.com",
@@ -42,7 +54,14 @@ function makeDeps(overrides: {
     },
   };
 
-  return { deps, parseSpy, enrichFromApiSpy, getApiKeySpy };
+  return {
+    deps,
+    parseSpy,
+    enrichFromApiSpy,
+    getApiKeySpy,
+    findImportsSpy,
+    getWrappersSpy,
+  };
 }
 
 describe("enrichFileForAgent", () => {
@@ -97,8 +116,8 @@ describe("enrichFileForAgent", () => {
     expect(enrichFromApiSpy).not.toHaveBeenCalled();
   });
 
-  test("returns null and skips parse when content has no posthog reference", async () => {
-    const { deps, parseSpy } = makeDeps({});
+  test("returns null and skips parse when content has no posthog reference AND no relative imports", async () => {
+    const { deps, parseSpy, findImportsSpy } = makeDeps({});
     const result = await enrichFileForAgent(
       deps,
       "/tmp/code.ts",
@@ -106,6 +125,121 @@ describe("enrichFileForAgent", () => {
     );
     expect(result).toBeNull();
     expect(parseSpy).not.toHaveBeenCalled();
+    expect(findImportsSpy).not.toHaveBeenCalled();
+  });
+
+  test("relative import with no resolvable wrapper → skips parse", async () => {
+    const { deps, parseSpy, findImportsSpy, getWrappersSpy } = makeDeps({
+      findImportsInSource: async () => [
+        {
+          localName: "foo",
+          importedName: "foo",
+          resolvedAbsPath: "/tmp/foo.ts",
+        },
+      ],
+      getWrappersForFile: async () => [],
+    });
+    const result = await enrichFileForAgent(
+      deps,
+      "/tmp/app.ts",
+      'import { foo } from "./foo";\nfoo("x");',
+    );
+    expect(result).toBeNull();
+    expect(findImportsSpy).toHaveBeenCalled();
+    expect(getWrappersSpy).toHaveBeenCalledWith("/tmp/foo.ts");
+    expect(parseSpy).not.toHaveBeenCalled();
+  });
+
+  test("relative import hitting a named wrapper triggers parse with context", async () => {
+    const wrapper = {
+      name: "track",
+      methodKind: "capture",
+      posthogMethod: "capture",
+      classification: { kind: "pass-through", paramIndex: 0 },
+      isNamedExport: true,
+      isDefaultExport: false,
+    };
+    const { deps, parseSpy, findImportsSpy, getWrappersSpy } = makeDeps({
+      findImportsInSource: async () => [
+        {
+          localName: "track",
+          importedName: "track",
+          resolvedAbsPath: "/tmp/telemetry.ts",
+        },
+      ],
+      getWrappersForFile: async () => [wrapper],
+    });
+    const result = await enrichFileForAgent(
+      deps,
+      "/tmp/app.ts",
+      'import { track } from "./telemetry";\ntrack("x");',
+    );
+    expect(result).toBe("enriched content");
+    expect(findImportsSpy).toHaveBeenCalled();
+    expect(getWrappersSpy).toHaveBeenCalledWith("/tmp/telemetry.ts");
+    expect(parseSpy).toHaveBeenCalledWith(
+      expect.any(String),
+      expect.any(String),
+      expect.objectContaining({
+        wrappersByLocalName: expect.any(Map),
+      }),
+    );
+    const ctxArg = parseSpy.mock.calls[0][2] as {
+      wrappersByLocalName: Map<string, unknown>;
+    };
+    expect(ctxArg.wrappersByLocalName.get("track")).toEqual(wrapper);
+  });
+
+  test("file with posthog literal and no relative imports skips import resolution", async () => {
+    const { deps, findImportsSpy, parseSpy } = makeDeps({});
+    await enrichFileForAgent(deps, "/tmp/code.ts", "posthog.capture('x');");
+    expect(findImportsSpy).not.toHaveBeenCalled();
+    expect(parseSpy).toHaveBeenCalled();
+  });
+
+  test("file with only bare-package imports does not trigger import resolution", async () => {
+    const { deps, findImportsSpy } = makeDeps({});
+    const content = [
+      'import React from "react";',
+      'import { useState } from "react";',
+      'import posthog from "posthog-js";',
+      "posthog.capture('x');",
+    ].join("\n");
+    await enrichFileForAgent(deps, "/tmp/page.tsx", content);
+    expect(findImportsSpy).not.toHaveBeenCalled();
+  });
+
+  test("file with direct posthog AND wrapper imports gets both enriched", async () => {
+    const wrapper = {
+      name: "track",
+      methodKind: "capture",
+      posthogMethod: "capture",
+      classification: { kind: "pass-through", paramIndex: 0 },
+      isNamedExport: true,
+      isDefaultExport: false,
+    };
+    const { deps, findImportsSpy, parseSpy } = makeDeps({
+      findImportsInSource: async () => [
+        {
+          localName: "track",
+          importedName: "track",
+          resolvedAbsPath: "/tmp/utils.ts",
+        },
+      ],
+      getWrappersForFile: async () => [wrapper],
+    });
+    const content = [
+      'import posthog from "posthog-js";',
+      'import { track } from "./utils";',
+      "posthog.capture('direct');",
+      'track("wrapper_call");',
+    ].join("\n");
+    await enrichFileForAgent(deps, "/tmp/page.tsx", content);
+    expect(findImportsSpy).toHaveBeenCalled();
+    const ctx = parseSpy.mock.calls[0][2] as {
+      wrappersByLocalName: Map<string, unknown>;
+    };
+    expect(ctx.wrappersByLocalName.get("track")).toEqual(wrapper);
   });
 
   test("returns null when getApiKey yields empty string", async () => {

package/src/enrichment/file-enricher.ts

@@ -1,5 +1,11 @@
 import * as path from "node:path";
-import { EXT_TO_LANG_ID, PostHogEnricher } from "@posthog/enricher";
+import {
+  EXT_TO_LANG_ID,
+  type ImportEdge,
+  type LocalWrapper,
+  type ParseContext,
+  PostHogEnricher,
+} from "@posthog/enricher";
 import type { PostHogAPIConfig } from "../types";
 import type { Logger } from "../utils/logger";
 
@@ -27,6 +33,10 @@ export function createEnrichment(
 }
 
 const MAX_ENRICHMENT_BYTES = 1_000_000;
+const MAX_RELATIVE_IMPORTS = 64;
+const RELATIVE_IMPORT_REGEX =
+  /(?:^|\n)\s*(?:import\b[^\n]*['"]\.{1,2}\/|from\s+\.)/;
+const POSTHOG_LITERAL_REGEX = /posthog/i;
 
 export async function enrichFileForAgent(
   deps: FileEnrichmentDeps,
@@ -35,15 +45,29 @@ export async function enrichFileForAgent(
 ): Promise<string | null> {
   if (!content || content.length > MAX_ENRICHMENT_BYTES) return null;
 
-  // Skip the tree-sitter parse for files with no PostHog references.
-  if (!/posthog/i.test(content)) return null;
-
   const ext = path.extname(filePath).toLowerCase();
   const langId = EXT_TO_LANG_ID[ext];
   if (!langId || !deps.enricher.isSupported(langId)) return null;
 
+  const hasPostHogLiteral = POSTHOG_LITERAL_REGEX.test(content);
+  const hasRelativeImport = RELATIVE_IMPORT_REGEX.test(content);
+  let parseContext: ParseContext | undefined;
+
+  // Build wrapper context whenever the file has relative imports — direct PostHog
+  // usage and wrapper usage can coexist in the same file, so we don't skip this
+  // just because `posthog` already appears literally.
+  if (hasRelativeImport) {
+    const absPath = path.resolve(filePath);
+    const ctx = await buildWrapperContext(deps, content, langId, absPath);
+    if (ctx) parseContext = ctx;
+  }
+
+  // Bail only when nothing at all could be enriched: no direct posthog literal
+  // AND no resolvable wrappers.
+  if (!hasPostHogLiteral && !parseContext) return null;
+
   try {
-    const parsed = await deps.enricher.parse(content, langId);
+    const parsed = await deps.enricher.parse(content, langId, parseContext);
     if (parsed.calls.length === 0 && parsed.initCalls.length === 0) {
       return null;
     }
@@ -69,6 +93,7 @@ export async function enrichFileForAgent(
     deps.logger?.debug("File enriched", {
       filePath,
       calls: parsed.calls.length,
+      viaWrappers: parsed.calls.filter((c) => c.viaWrapper).length,
     });
     return annotated;
   } catch (err) {
@@ -80,3 +105,69 @@ export async function enrichFileForAgent(
     return null;
   }
 }
+
+async function buildWrapperContext(
+  deps: FileEnrichmentDeps,
+  content: string,
+  langId: string,
+  absPath: string,
+): Promise<ParseContext | null> {
+  let edges: ImportEdge[];
+  try {
+    edges = await deps.enricher.findImportsInSource(content, langId, absPath);
+  } catch (err) {
+    deps.logger?.debug("Import resolution failed", {
+      absPath,
+      err: err instanceof Error ? err.message : String(err),
+    });
+    return null;
+  }
+
+  if (!edges.length) return null;
+  const bounded = edges.slice(0, MAX_RELATIVE_IMPORTS);
+
+  const wrappersByLocalName = new Map<string, LocalWrapper>();
+  const namespaceWrappers = new Map<string, Map<string, LocalWrapper>>();
+
+  const resolutions = await Promise.all(
+    bounded.map(async (edge) => {
+      if (!edge.resolvedAbsPath) return null;
+      const wrappers = await deps.enricher.getWrappersForFile(
+        edge.resolvedAbsPath,
+      );
+      if (!wrappers.length) return null;
+      return { edge, wrappers };
+    }),
+  );
+
+  for (const entry of resolutions) {
+    if (!entry) continue;
+    const { edge, wrappers } = entry;
+
+    if (edge.isNamespace) {
+      const nsMap = new Map<string, LocalWrapper>();
+      for (const w of wrappers) {
+        if (w.isNamedExport || w.isDefaultExport) {
+          nsMap.set(w.name, w);
+        }
+      }
+      if (nsMap.size) namespaceWrappers.set(edge.localName, nsMap);
+      continue;
+    }
+
+    if (edge.isDefault) {
+      const target = wrappers.find((w) => w.isDefaultExport);
+      if (target) wrappersByLocalName.set(edge.localName, target);
+      continue;
+    }
+
+    const target = wrappers.find(
+      (w) => w.name === edge.importedName && w.isNamedExport,
+    );
+    if (target) wrappersByLocalName.set(edge.localName, target);
+  }
+
+  if (!wrappersByLocalName.size && !namespaceWrappers.size) return null;
+
+  return { wrappersByLocalName, namespaceWrappers };
+}