diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index a9eed0d58..2aff79fdf 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -42,6 +42,9 @@ jobs: - name: Build git run: pnpm --filter @posthog/git build + - name: Build enricher + run: pnpm --filter @posthog/enricher build + - name: Build agent run: pnpm --filter agent build diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index db5b51cbe..6c0f92a07 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -81,6 +81,7 @@ jobs: pnpm --filter @posthog/platform build & pnpm --filter @posthog/shared build pnpm --filter @posthog/git build + pnpm --filter @posthog/enricher build pnpm --filter agent build & wait diff --git a/apps/code/forge.config.ts b/apps/code/forge.config.ts index ea6f4c240..ac1506cdf 100644 --- a/apps/code/forge.config.ts +++ b/apps/code/forge.config.ts @@ -143,7 +143,7 @@ const config: ForgeConfig = { packagerConfig: { asar: { unpack: - "{**/*.node,**/spawn-helper,**/.vite/build/claude-cli/**,**/.vite/build/plugins/posthog/**,**/.vite/build/codex-acp/**,**/node_modules/node-pty/**,**/node_modules/@parcel/**,**/node_modules/file-icon/**,**/node_modules/better-sqlite3/**,**/node_modules/bindings/**,**/node_modules/file-uri-to-path/**}", + "{**/*.node,**/spawn-helper,**/.vite/build/claude-cli/**,**/.vite/build/plugins/posthog/**,**/.vite/build/codex-acp/**,**/.vite/build/grammars/**,**/node_modules/node-pty/**,**/node_modules/@parcel/**,**/node_modules/file-icon/**,**/node_modules/better-sqlite3/**,**/node_modules/bindings/**,**/node_modules/file-uri-to-path/**}", }, prune: false, name: "PostHog Code", diff --git a/apps/code/vite.main.config.mts b/apps/code/vite.main.config.mts index 2586f1013..402b6c341 100644 --- a/apps/code/vite.main.config.mts +++ b/apps/code/vite.main.config.mts @@ -424,6 +424,52 @@ function copyDrizzleMigrations(): Plugin { }; } +let enricherGrammarsCopied = false; + +function copyEnricherGrammars(): Plugin { + return { + name: "copy-enricher-grammars", + writeBundle() { + // `.vite/grammars` is what the bundle resolves at dev-time; electron-forge + // only copies `.vite/build/**` into the packaged app, so we need both. + const destDirs = [ + join(__dirname, ".vite/grammars"), + join(__dirname, ".vite/build/grammars"), + ]; + + if (enricherGrammarsCopied && destDirs.every((d) => existsSync(d))) { + return; + } + + const candidates = [ + join(__dirname, "node_modules/@posthog/enricher/grammars"), + join(__dirname, "../../node_modules/@posthog/enricher/grammars"), + join(__dirname, "../../packages/enricher/grammars"), + ]; + + const sourceDir = candidates.find((p) => existsSync(p)); + if (!sourceDir) { + console.warn( + "[copy-enricher-grammars] grammars directory not found. Checked:", + candidates.join(", "), + ); + return; + } + + for (const destDir of destDirs) { + if (!existsSync(destDir)) { + mkdirSync(destDir, { recursive: true }); + } + cpSync(sourceDir, destDir, { recursive: true }); + } + enricherGrammarsCopied = true; + console.log( + `Copied enricher grammars from ${sourceDir} to ${destDirs.join(", ")}`, + ); + }, + }; +} + let codexAcpCopied = false; function copyCodexAcpBinaries(): Plugin { @@ -492,6 +538,7 @@ export default defineConfig(({ mode }) => { copyPosthogPlugin(isDev), copyDrizzleMigrations(), copyCodexAcpBinaries(), + copyEnricherGrammars(), createPosthogPlugin(env, "posthog-code-main"), ].filter(Boolean), define: { diff --git a/packages/agent/package.json b/packages/agent/package.json index a60936427..f6f10d63f 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -109,6 +109,7 @@ "@opentelemetry/resources": "^2.0.0", "@opentelemetry/sdk-logs": "^0.208.0", "@opentelemetry/semantic-conventions": "^1.28.0", + "@posthog/enricher": "workspace:*", "@types/jsonwebtoken": "^9.0.10", "commander": "^14.0.2", "hono": "^4.11.7", diff --git a/packages/agent/src/adapters/acp-connection.ts b/packages/agent/src/adapters/acp-connection.ts index cbd834cd1..085ca58be 100644 --- a/packages/agent/src/adapters/acp-connection.ts +++ b/packages/agent/src/adapters/acp-connection.ts @@ -1,6 +1,6 @@ import { AgentSideConnection, ndJsonStream } from "@agentclientprotocol/sdk"; import type { SessionLogWriter } from "../session-log-writer"; -import type { ProcessSpawnedCallback } from "../types"; +import type { PostHogAPIConfig, ProcessSpawnedCallback } from "../types"; import { Logger } from "../utils/logger"; import { createBidirectionalStreams, @@ -26,6 +26,10 @@ export type AcpConnectionConfig = { allowedModelIds?: Set; /** Callback invoked when the agent calls the create_output tool for structured output */ onStructuredOutput?: (output: Record) => Promise; + /** PostHog API config; when set, enables file-read enrichment unless disabled. */ + posthogApiConfig?: PostHogAPIConfig; + /** Defaults to true when posthogApiConfig is set. Set to false to disable enrichment. */ + enricherEnabled?: boolean; }; export type AcpConnection = { @@ -54,6 +58,13 @@ export function createAcpConnection( return createClaudeConnection(config); } +function resolveEnricherApiConfig( + config: AcpConnectionConfig, +): PostHogAPIConfig | undefined { + const enabled = !!config.posthogApiConfig && config.enricherEnabled !== false; + return enabled ? config.posthogApiConfig : undefined; +} + function createClaudeConnection(config: AcpConnectionConfig): AcpConnection { const logger = config.logger?.child("AcpConnection") ?? @@ -102,6 +113,7 @@ function createClaudeConnection(config: AcpConnectionConfig): AcpConnection { agent = new ClaudeAcpAgent(client, { ...config.processCallbacks, onStructuredOutput: config.onStructuredOutput, + posthogApiConfig: resolveEnricherApiConfig(config), }); return agent; }, agentStream); @@ -192,6 +204,7 @@ function createCodexConnection(config: AcpConnectionConfig): AcpConnection { agent = new CodexAcpAgent(client, { codexProcessOptions: config.codexOptions ?? {}, processCallbacks: config.processCallbacks, + posthogApiConfig: resolveEnricherApiConfig(config), }); return agent; }, agentStream); diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index b81736e62..d7d0e30bc 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -51,6 +51,12 @@ import { POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "../../acp-extensions"; +import { + createEnrichment, + type Enrichment, + type FileEnrichmentDeps, +} from "../../enrichment/file-enricher"; +import type { PostHogAPIConfig } from "../../types"; import { unreachable, withTimeout } from "../../utils/common"; import { Logger } from "../../utils/logger"; import { Pushable } from "../../utils/streams"; @@ -62,6 +68,7 @@ import { handleSystemMessage, handleUserAssistantMessage, } from "./conversion/sdk-to-acp"; +import type { EnrichedReadCache } from "./hooks"; import { fetchMcpToolMetadata, getConnectedMcpServerNames, @@ -116,6 +123,7 @@ export interface ClaudeAcpAgentOptions { onProcessExited?: (pid: number) => void; onMcpServersReady?: (serverNames: string[]) => void; onStructuredOutput?: (output: Record) => Promise; + posthogApiConfig?: PostHogAPIConfig; } export class ClaudeAcpAgent extends BaseAcpAgent { @@ -125,12 +133,29 @@ export class ClaudeAcpAgent extends BaseAcpAgent { backgroundTerminals: { [key: string]: BackgroundTerminal } = {}; clientCapabilities?: ClientCapabilities; private options?: ClaudeAcpAgentOptions; + private enrichment?: Enrichment; + private enrichedReadCache: EnrichedReadCache = new Map(); constructor(client: AgentSideConnection, options?: ClaudeAcpAgentOptions) { super(client); this.options = options; this.toolUseCache = {}; this.logger = new Logger({ debug: true, prefix: "[ClaudeAcpAgent]" }); + this.enrichment = createEnrichment(options?.posthogApiConfig, this.logger); + } + + protected getEnrichmentDeps(): FileEnrichmentDeps | undefined { + return this.enrichment?.deps; + } + + override async closeSession(): Promise { + try { + await super.closeSession(); + } finally { + this.enrichment?.dispose(); + this.enrichment = undefined; + this.enrichedReadCache.clear(); + } } async initialize(request: InitializeRequest): Promise { @@ -355,6 +380,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { client: this.client, toolUseCache: this.toolUseCache, fileContentCache: this.fileContentCache, + enrichedReadCache: this.enrichedReadCache, logger: this.logger, supportsTerminalOutput, }; @@ -993,6 +1019,8 @@ export class ClaudeAcpAgent extends BaseAcpAgent { onProcessSpawned: this.options?.onProcessSpawned, onProcessExited: this.options?.onProcessExited, effort, + enrichmentDeps: this.enrichment?.deps, + enrichedReadCache: this.enrichedReadCache, }); // Use the same abort controller that buildSessionOptions gave to the query @@ -1354,6 +1382,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { client: this.client, toolUseCache: this.toolUseCache, fileContentCache: this.fileContentCache, + enrichedReadCache: this.enrichedReadCache, logger: this.logger, registerHooks: false, }; diff --git a/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts b/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts index cf0e7b4d3..ff6e55a02 100644 --- a/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts +++ b/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts @@ -21,7 +21,7 @@ import { POSTHOG_NOTIFICATIONS } from "@/acp-extensions"; import { image, text } from "../../../utils/acp-content"; import { unreachable } from "../../../utils/common"; import type { Logger } from "../../../utils/logger"; -import { registerHookCallback } from "../hooks"; +import { type EnrichedReadCache, registerHookCallback } from "../hooks"; import type { Session, ToolUpdateMeta, ToolUseCache } from "../types"; import { type ClaudePlanEntry, @@ -51,6 +51,7 @@ type ChunkHandlerContext = { sessionId: string; toolUseCache: ToolUseCache; fileContentCache: { [key: string]: string }; + enrichedReadCache?: EnrichedReadCache; client: AgentSideConnection; logger: Logger; parentToolCallId?: string; @@ -67,6 +68,7 @@ export interface MessageHandlerContext { client: AgentSideConnection; toolUseCache: ToolUseCache; fileContentCache: { [key: string]: string }; + enrichedReadCache?: EnrichedReadCache; logger: Logger; registerHooks?: boolean; supportsTerminalOutput?: boolean; @@ -248,7 +250,7 @@ function extractTextFromContent(content: unknown): string | null { return null; } -function stripCatLineNumbers(text: string): string { +export function stripCatLineNumbers(text: string): string { return text.replace(/^ *\d+[\t→]/gm, ""); } @@ -318,6 +320,7 @@ function handleToolResultChunk( supportsTerminalOutput: ctx.supportsTerminalOutput, toolUseId: chunk.tool_use_id, cachedFileContent: ctx.fileContentCache, + enrichedReadCache: ctx.enrichedReadCache, }, ); @@ -448,6 +451,7 @@ function toAcpNotifications( supportsTerminalOutput?: boolean, cwd?: string, mcpToolUseResult?: Record, + enrichedReadCache?: EnrichedReadCache, ): SessionNotification[] { if (typeof content === "string") { const update: SessionUpdate = { @@ -468,6 +472,7 @@ function toAcpNotifications( sessionId, toolUseCache, fileContentCache, + enrichedReadCache, client, logger, parentToolCallId, @@ -498,6 +503,7 @@ function streamEventToAcpNotifications( registerHooks?: boolean, supportsTerminalOutput?: boolean, cwd?: string, + enrichedReadCache?: EnrichedReadCache, ): SessionNotification[] { const event = message.event; switch (event.type) { @@ -514,6 +520,8 @@ function streamEventToAcpNotifications( registerHooks, supportsTerminalOutput, cwd, + undefined, + enrichedReadCache, ); case "content_block_delta": return toAcpNotifications( @@ -528,6 +536,8 @@ function streamEventToAcpNotifications( registerHooks, supportsTerminalOutput, cwd, + undefined, + enrichedReadCache, ); case "message_start": case "message_delta": @@ -717,6 +727,7 @@ export async function handleStreamEvent( context.registerHooks, context.supportsTerminalOutput, context.session.cwd, + context.enrichedReadCache, )) { await client.sessionUpdate(notification); context.session.notificationHistory.push(notification); @@ -840,6 +851,7 @@ export async function handleUserAssistantMessage( context.supportsTerminalOutput, session.cwd, mcpToolUseResult, + context.enrichedReadCache, )) { await client.sessionUpdate(notification); session.notificationHistory.push(notification); diff --git a/packages/agent/src/adapters/claude/conversion/tool-use-to-acp.ts b/packages/agent/src/adapters/claude/conversion/tool-use-to-acp.ts index c66017a3f..e960bd467 100644 --- a/packages/agent/src/adapters/claude/conversion/tool-use-to-acp.ts +++ b/packages/agent/src/adapters/claude/conversion/tool-use-to-acp.ts @@ -31,6 +31,7 @@ function stripSystemReminders(value: string): string { } import { resourceLink, text, toolContent } from "../../../utils/acp-content"; +import type { EnrichedReadCache } from "../hooks"; import { getMcpToolMetadata } from "../mcp/tool-metadata"; type ToolInfo = Pick; @@ -526,6 +527,7 @@ export function toolUpdateFromToolResult( supportsTerminalOutput?: boolean; toolUseId?: string; cachedFileContent?: Record; + enrichedReadCache?: EnrichedReadCache; }, ): Pick { if ( @@ -538,7 +540,21 @@ export function toolUpdateFromToolResult( } switch (toolUse?.name) { - case "Read": + case "Read": { + const cache = options?.enrichedReadCache; + const enriched = + cache && options?.toolUseId ? cache.get(options.toolUseId) : undefined; + if (enriched !== undefined && cache && options?.toolUseId) { + cache.delete(options.toolUseId); + return { + content: [ + { + type: "content" as const, + content: text(markdownEscape(enriched)), + }, + ], + }; + } if (Array.isArray(toolResult.content) && toolResult.content.length > 0) { return { content: toolResult.content.map((item) => { @@ -582,6 +598,7 @@ export function toolUpdateFromToolResult( }; } return {}; + } case "Bash": { const result = toolResult.content; diff --git a/packages/agent/src/adapters/claude/hooks.test.ts b/packages/agent/src/adapters/claude/hooks.test.ts new file mode 100644 index 000000000..ec096d79d --- /dev/null +++ b/packages/agent/src/adapters/claude/hooks.test.ts @@ -0,0 +1,189 @@ +import type { HookInput } from "@anthropic-ai/claude-agent-sdk"; +import { describe, expect, test, vi } from "vitest"; +import type { FileEnrichmentDeps } from "../../enrichment/file-enricher"; + +const enrichFileMock = vi.hoisted(() => vi.fn()); +vi.mock("../../enrichment/file-enricher", () => ({ + enrichFileForAgent: enrichFileMock, +})); + +import { createReadEnrichmentHook, type EnrichedReadCache } from "./hooks"; + +const stubDeps = {} as FileEnrichmentDeps; + +function buildReadHookInput( + overrides: Partial & { + file_path?: string; + tool_response?: unknown; + } = {}, +): HookInput { + return { + session_id: "test-session", + transcript_path: "/tmp/transcript", + cwd: "/tmp", + hook_event_name: "PostToolUse", + tool_name: "Read", + tool_use_id: "toolu_1", + tool_input: { file_path: overrides.file_path ?? "/tmp/code.ts" }, + tool_response: overrides.tool_response ?? "raw-content", + ...overrides, + } as HookInput; +} + +describe("createReadEnrichmentHook", () => { + test("returns { continue: true } for non-PostToolUse events", async () => { + enrichFileMock.mockReset(); + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + const result = await hook( + { hook_event_name: "PreToolUse" } as HookInput, + undefined, + { signal: new AbortController().signal }, + ); + expect(result).toEqual({ continue: true }); + expect(enrichFileMock).not.toHaveBeenCalled(); + }); + + test("returns { continue: true } for non-Read tools", async () => { + enrichFileMock.mockReset(); + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + const result = await hook( + buildReadHookInput({ tool_name: "Bash" }), + undefined, + { signal: new AbortController().signal }, + ); + expect(result).toEqual({ continue: true }); + expect(enrichFileMock).not.toHaveBeenCalled(); + }); + + test("passes stripped content and file_path into enricher", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce(null); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + await hook( + buildReadHookInput({ + file_path: "/tmp/app.ts", + tool_response: " 1\tconst x = 1;\n 2\tposthog.capture('x');", + }), + undefined, + { signal: new AbortController().signal }, + ); + + expect(enrichFileMock).toHaveBeenCalledTimes(1); + const [, filePath, content] = enrichFileMock.mock.calls[0]; + expect(filePath).toBe("/tmp/app.ts"); + expect(content).toBe("const x = 1;\nposthog.capture('x');"); + }); + + test("returns additionalContext when enricher produces annotations", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce( + "posthog.capture('x'); // [PostHog] Event: \"x\"", + ); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + const result = await hook( + buildReadHookInput({ file_path: "/tmp/app.ts" }), + undefined, + { + signal: new AbortController().signal, + }, + ); + + expect(result).toEqual({ + continue: true, + hookSpecificOutput: { + hookEventName: "PostToolUse", + additionalContext: expect.stringContaining( + "posthog.capture('x'); // [PostHog] Event: \"x\"", + ), + }, + }); + const context = ( + result as { + hookSpecificOutput: { additionalContext: string }; + } + ).hookSpecificOutput.additionalContext; + expect(context).toContain("/tmp/app.ts"); + }); + + test("writes enriched content to cache keyed by tool_use_id", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce( + "posthog.capture('x'); // [PostHog] Event: \"x\"", + ); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + await hook(buildReadHookInput({ file_path: "/tmp/app.ts" }), undefined, { + signal: new AbortController().signal, + }); + + expect(cache.get("toolu_1")).toContain('// [PostHog] Event: "x"'); + }); + + test("does not write to cache when tool_use_id is missing", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce("enriched"); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + await hook( + buildReadHookInput({ file_path: "/tmp/app.ts", tool_use_id: undefined }), + undefined, + { signal: new AbortController().signal }, + ); + + expect(cache.size).toBe(0); + }); + + test("handles {type:'text', file:{content}} Read tool_response shape", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce("enriched"); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + await hook( + buildReadHookInput({ + file_path: "/tmp/app.ts", + tool_response: { + type: "text", + file: { + filePath: "/tmp/app.ts", + content: "posthog.capture('x');\n", + numLines: 1, + startLine: 1, + totalLines: 1, + }, + }, + }), + undefined, + { signal: new AbortController().signal }, + ); + + const [, , content] = enrichFileMock.mock.calls[0]; + expect(content).toBe("posthog.capture('x');\n"); + }); + + test("handles wrapped [{type:'text', text:'...'}] tool_response shape", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce("enriched"); + + const cache: EnrichedReadCache = new Map(); + const hook = createReadEnrichmentHook(stubDeps, cache); + await hook( + buildReadHookInput({ + tool_response: [{ type: "text", text: " 1\tfoo" }], + }), + undefined, + { signal: new AbortController().signal }, + ); + + const [, , content] = enrichFileMock.mock.calls[0]; + expect(content).toBe("foo"); + }); +}); diff --git a/packages/agent/src/adapters/claude/hooks.ts b/packages/agent/src/adapters/claude/hooks.ts index c85028c55..3d704d5ff 100644 --- a/packages/agent/src/adapters/claude/hooks.ts +++ b/packages/agent/src/adapters/claude/hooks.ts @@ -1,8 +1,101 @@ import type { HookCallback, HookInput } from "@anthropic-ai/claude-agent-sdk"; +import { + enrichFileForAgent, + type FileEnrichmentDeps, +} from "../../enrichment/file-enricher"; import type { Logger } from "../../utils/logger"; +import { stripCatLineNumbers } from "./conversion/sdk-to-acp"; import type { SettingsManager } from "./session/settings"; import type { CodeExecutionMode } from "./tools"; +function extractTextFromToolResponse(response: unknown): string | null { + if (typeof response === "string") return response; + if (!response) return null; + if (Array.isArray(response)) { + const parts: string[] = []; + for (const part of response) { + if (typeof part === "string") { + parts.push(part); + } else if ( + part && + typeof part === "object" && + "text" in part && + typeof (part as { text?: unknown }).text === "string" + ) { + parts.push((part as { text: string }).text); + } + } + return parts.length > 0 ? parts.join("") : null; + } + if (typeof response === "object" && response !== null) { + const maybe = response as { + content?: unknown; + text?: unknown; + file?: { content?: unknown }; + }; + if ( + maybe.file && + typeof maybe.file === "object" && + typeof maybe.file.content === "string" + ) { + return maybe.file.content; + } + if (typeof maybe.text === "string") return maybe.text; + if (maybe.content) return extractTextFromToolResponse(maybe.content); + } + return null; +} + +/** + * Per-toolUseId handoff from the PostToolUse hook to `toolUpdateFromToolResult`. + * Can't emit a standalone `tool_call_update` because the SDK emits its own + * when it processes the tool_result, and the renderer applies it via + * `Object.assign` — our earlier update would be overwritten. + */ +export type EnrichedReadCache = Map; + +export const createReadEnrichmentHook = + (deps: FileEnrichmentDeps, cache: EnrichedReadCache): HookCallback => + async (input: HookInput) => { + if (input.hook_event_name !== "PostToolUse") return { continue: true }; + if (input.tool_name !== "Read") return { continue: true }; + + const toolInput = input.tool_input as { file_path?: string } | undefined; + const filePath = toolInput?.file_path; + if (!filePath) return { continue: true }; + + const raw = extractTextFromToolResponse(input.tool_response); + if (!raw) return { continue: true }; + + const enriched = await enrichFileForAgent( + deps, + filePath, + stripCatLineNumbers(raw), + ); + if (!enriched) return { continue: true }; + + if (input.tool_use_id) { + cache.set(input.tool_use_id, enriched); + } + + return { + continue: true, + hookSpecificOutput: { + hookEventName: "PostToolUse" as const, + additionalContext: [ + `## PostHog metadata for ${filePath}`, + "", + "The file below is annotated with live data from the user's PostHog project:", + "flag type / rollout / staleness / linked experiment, and for events the verification status,", + "30-day volume, and unique-user count. Treat these as authoritative product context —", + "they describe what is actually running in production.", + "", + enriched, + ].join("\n"), + }, + }; + }; + const toolUseCallbacks: { [toolUseId: string]: { onPostToolUseHook?: ( diff --git a/packages/agent/src/adapters/claude/session/options.ts b/packages/agent/src/adapters/claude/session/options.ts index 03506baf2..0f9d80186 100644 --- a/packages/agent/src/adapters/claude/session/options.ts +++ b/packages/agent/src/adapters/claude/session/options.ts @@ -10,12 +10,15 @@ import type { SpawnedProcess, SpawnOptions, } from "@anthropic-ai/claude-agent-sdk"; +import type { FileEnrichmentDeps } from "../../../enrichment/file-enricher"; import { IS_ROOT } from "../../../utils/common"; import type { Logger } from "../../../utils/logger"; import { createPostToolUseHook, createPreToolUseHook, + createReadEnrichmentHook, createSubagentRewriteHook, + type EnrichedReadCache, type OnModeChange, } from "../hooks"; import type { CodeExecutionMode } from "../tools"; @@ -49,6 +52,8 @@ export interface BuildOptionsParams { onProcessSpawned?: (info: ProcessSpawnedInfo) => void; onProcessExited?: (pid: number) => void; effort?: EffortLevel; + enrichmentDeps?: FileEnrichmentDeps; + enrichedReadCache?: EnrichedReadCache; } export function buildSystemPrompt( @@ -108,14 +113,21 @@ function buildHooks( onModeChange: OnModeChange | undefined, settingsManager: SettingsManager, logger: Logger, + enrichmentDeps: FileEnrichmentDeps | undefined, + enrichedReadCache: EnrichedReadCache | undefined, ): Options["hooks"] { + const postToolUseHooks = [createPostToolUseHook({ onModeChange, logger })]; + if (enrichmentDeps && enrichedReadCache) { + postToolUseHooks.push( + createReadEnrichmentHook(enrichmentDeps, enrichedReadCache), + ); + } + return { ...userHooks, PostToolUse: [ ...(userHooks?.PostToolUse || []), - { - hooks: [createPostToolUseHook({ onModeChange, logger })], - }, + { hooks: postToolUseHooks }, ], PreToolUse: [ ...(userHooks?.PreToolUse || []), @@ -269,6 +281,8 @@ export function buildSessionOptions(params: BuildOptionsParams): Options { params.onModeChange, params.settingsManager, params.logger, + params.enrichmentDeps, + params.enrichedReadCache, ), outputFormat: params.outputFormat, abortController: getAbortController( diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index e73882151..d7a542bd3 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -41,6 +41,10 @@ import { POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "../../acp-extensions"; +import { + createEnrichment, + type Enrichment, +} from "../../enrichment/file-enricher"; import { type CodeExecutionMode, type CodexNativeMode, @@ -48,7 +52,7 @@ import { isCodexNativeMode, type PermissionMode, } from "../../execution-mode"; -import type { ProcessSpawnedCallback } from "../../types"; +import type { PostHogAPIConfig, ProcessSpawnedCallback } from "../../types"; import { Logger } from "../../utils/logger"; import { nodeReadableToWebReadable, @@ -86,6 +90,7 @@ interface NewSessionMeta { export interface CodexAcpAgentOptions { codexProcessOptions: CodexProcessOptions; processCallbacks?: ProcessSpawnedCallback; + posthogApiConfig?: PostHogAPIConfig; } type CodexSession = BaseSession & { @@ -168,6 +173,7 @@ export class CodexAcpAgent extends BaseAcpAgent { // Snapshot of the initialize() request so refreshSession can replay the // same handshake against a respawned codex-acp subprocess. private lastInitRequest?: InitializeRequest; + private enrichment?: Enrichment; constructor(client: AgentSideConnection, options: CodexAcpAgentOptions) { super(client); @@ -205,12 +211,16 @@ export class CodexAcpAgent extends BaseAcpAgent { this.sessionState = createSessionState("", cwd); + this.enrichment = createEnrichment(options.posthogApiConfig, this.logger); + // Create the ClientSideConnection to codex-acp. // The Client handler delegates all requests from codex-acp to the upstream // PostHog Code client via our AgentSideConnection. this.codexConnection = new ClientSideConnection( (_agent) => - createCodexClient(this.client, this.logger, this.sessionState), + createCodexClient(this.client, this.logger, this.sessionState, { + enrichmentDeps: this.enrichment?.deps, + }), codexStream, ); } @@ -672,5 +682,7 @@ export class CodexAcpAgent extends BaseAcpAgent { } catch (err) { this.logger.warn("Failed to kill codex-acp process", { error: err }); } + this.enrichment?.dispose(); + this.enrichment = undefined; } } diff --git a/packages/agent/src/adapters/codex/codex-client.test.ts b/packages/agent/src/adapters/codex/codex-client.test.ts new file mode 100644 index 000000000..c012a2b4b --- /dev/null +++ b/packages/agent/src/adapters/codex/codex-client.test.ts @@ -0,0 +1,112 @@ +import type { + AgentSideConnection, + ReadTextFileRequest, + ReadTextFileResponse, +} from "@agentclientprotocol/sdk"; +import { describe, expect, test, vi } from "vitest"; +import type { FileEnrichmentDeps } from "../../enrichment/file-enricher"; +import { Logger } from "../../utils/logger"; + +const enrichFileMock = vi.hoisted(() => vi.fn()); +vi.mock("../../enrichment/file-enricher", () => ({ + enrichFileForAgent: enrichFileMock, +})); + +import { createCodexClient } from "./codex-client"; +import { createSessionState } from "./session-state"; + +function makeUpstream(response: ReadTextFileResponse): AgentSideConnection & { + readTextFile: ReturnType; +} { + const mock = { + readTextFile: vi.fn(async (_: ReadTextFileRequest) => response), + writeTextFile: vi.fn(), + requestPermission: vi.fn(), + sessionUpdate: vi.fn(), + createTerminal: vi.fn(), + terminalOutput: vi.fn(), + releaseTerminal: vi.fn(), + waitForTerminalExit: vi.fn(), + killTerminal: vi.fn(), + extMethod: vi.fn(), + extNotification: vi.fn(), + }; + return mock as unknown as AgentSideConnection & { + readTextFile: ReturnType; + }; +} + +describe("createCodexClient readTextFile", () => { + const logger = new Logger({ debug: false, prefix: "[test]" }); + const sessionState = createSessionState("", "/tmp"); + + test("returns upstream response unchanged when enrichmentDeps is absent", async () => { + enrichFileMock.mockReset(); + const upstream = makeUpstream({ content: "const x = 1;" }); + const client = createCodexClient(upstream, logger, sessionState); + + const result = await client.readTextFile?.({ + sessionId: "s", + path: "/tmp/a.ts", + }); + expect(result?.content).toBe("const x = 1;"); + expect(enrichFileMock).not.toHaveBeenCalled(); + }); + + test("returns enriched content when helper returns a string", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce("const x = 1; // [PostHog] Flag ..."); + + const upstream = makeUpstream({ content: "const x = 1;" }); + const deps = {} as FileEnrichmentDeps; + const client = createCodexClient(upstream, logger, sessionState, { + enrichmentDeps: deps, + }); + + const result = await client.readTextFile?.({ + sessionId: "s", + path: "/tmp/a.ts", + }); + expect(result?.content).toBe("const x = 1; // [PostHog] Flag ..."); + expect(enrichFileMock).toHaveBeenCalledWith( + deps, + "/tmp/a.ts", + "const x = 1;", + ); + }); + + test("falls back to upstream response when helper returns null", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce(null); + + const upstream = makeUpstream({ content: "no posthog here" }); + const client = createCodexClient(upstream, logger, sessionState, { + enrichmentDeps: {} as FileEnrichmentDeps, + }); + + const result = await client.readTextFile?.({ + sessionId: "s", + path: "/tmp/a.ts", + }); + expect(result?.content).toBe("no posthog here"); + }); + + test("calls upstream.readTextFile with original params (UI sees original)", async () => { + enrichFileMock.mockReset(); + enrichFileMock.mockResolvedValueOnce("enriched"); + + const upstream = makeUpstream({ content: "original" }); + const client = createCodexClient(upstream, logger, sessionState, { + enrichmentDeps: {} as FileEnrichmentDeps, + }); + + const params = { + sessionId: "s", + path: "/tmp/a.ts", + line: 10, + limit: 5, + }; + await client.readTextFile?.(params); + expect(upstream.readTextFile).toHaveBeenCalledWith(params); + }); +}); diff --git a/packages/agent/src/adapters/codex/codex-client.ts b/packages/agent/src/adapters/codex/codex-client.ts index f56ab4493..727d8876b 100644 --- a/packages/agent/src/adapters/codex/codex-client.ts +++ b/packages/agent/src/adapters/codex/codex-client.ts @@ -29,6 +29,10 @@ import type { WriteTextFileRequest, WriteTextFileResponse, } from "@agentclientprotocol/sdk"; +import { + enrichFileForAgent, + type FileEnrichmentDeps, +} from "../../enrichment/file-enricher"; import type { PermissionMode } from "../../execution-mode"; import type { Logger } from "../../utils/logger"; import type { CodexSessionState } from "./session-state"; @@ -36,6 +40,8 @@ import type { CodexSessionState } from "./session-state"; export interface CodexClientCallbacks { /** Called when a usage_update session notification is received */ onUsageUpdate?: (update: Record) => void; + /** When set, Read responses are annotated with PostHog enrichment before reaching codex-acp. */ + enrichmentDeps?: FileEnrichmentDeps; } const AUTO_APPROVED_KINDS: Record> = { @@ -152,7 +158,14 @@ export function createCodexClient( async readTextFile( params: ReadTextFileRequest, ): Promise { - return upstreamClient.readTextFile(params); + const response = await upstreamClient.readTextFile(params); + if (!callbacks?.enrichmentDeps) return response; + const enriched = await enrichFileForAgent( + callbacks.enrichmentDeps, + params.path, + response.content, + ); + return enriched ? { ...response, content: enriched } : response; }, async writeTextFile( diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 522054468..f66b0fd74 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -19,6 +19,8 @@ export class Agent { private acpConnection?: InProcessAcpConnection; private taskRunId?: string; private sessionLogWriter?: SessionLogWriter; + private posthogApiConfig?: AgentConfig["posthog"]; + private enricherEnabled: boolean; constructor(config: AgentConfig) { this.logger = new Logger({ @@ -29,7 +31,9 @@ export class Agent { if (config.posthog) { this.posthogAPI = new PostHogAPIClient(config.posthog); + this.posthogApiConfig = config.posthog; } + this.enricherEnabled = config.enricher?.enabled !== false; if (config.posthog && !config.skipLogPersistence) { this.sessionLogWriter = new SessionLogWriter({ @@ -121,6 +125,8 @@ export class Agent { processCallbacks: options.processCallbacks, onStructuredOutput: options.onStructuredOutput, allowedModelIds, + posthogApiConfig: this.posthogApiConfig, + enricherEnabled: this.enricherEnabled, codexOptions: options.adapter === "codex" && gatewayConfig ? { diff --git a/packages/agent/src/enrichment/file-enricher.test.ts b/packages/agent/src/enrichment/file-enricher.test.ts new file mode 100644 index 000000000..751b70487 --- /dev/null +++ b/packages/agent/src/enrichment/file-enricher.test.ts @@ -0,0 +1,163 @@ +import { describe, expect, test, vi } from "vitest"; +import { enrichFileForAgent, type FileEnrichmentDeps } from "./file-enricher"; + +function makeDeps(overrides: { + toInlineCommentsReturn?: string; + callsCount?: number; + initCallsCount?: number; + parseRejects?: Error; + isSupported?: boolean; + getApiKey?: () => string | Promise; +}): { + deps: FileEnrichmentDeps; + parseSpy: ReturnType; + enrichFromApiSpy: ReturnType; + getApiKeySpy: ReturnType; +} { + const enrichFromApiSpy = vi.fn(async () => ({ + toInlineComments: () => + overrides.toInlineCommentsReturn ?? "enriched content", + })); + + const parseSpy = vi.fn(async () => { + if (overrides.parseRejects) throw overrides.parseRejects; + return { + calls: Array.from({ length: overrides.callsCount ?? 1 }), + initCalls: Array.from({ length: overrides.initCallsCount ?? 0 }), + enrichFromApi: enrichFromApiSpy, + }; + }); + + const getApiKeySpy = vi.fn(overrides.getApiKey ?? (() => "phx_test")); + + const deps: FileEnrichmentDeps = { + enricher: { + isSupported: vi.fn(() => overrides.isSupported ?? true), + parse: parseSpy, + } as unknown as FileEnrichmentDeps["enricher"], + apiConfig: { + apiUrl: "https://test.posthog.com", + projectId: 1, + getApiKey: getApiKeySpy, + }, + }; + + return { deps, parseSpy, enrichFromApiSpy, getApiKeySpy }; +} + +describe("enrichFileForAgent", () => { + test("returns null for unsupported extension", async () => { + const { deps, parseSpy } = makeDeps({}); + const result = await enrichFileForAgent( + deps, + "/tmp/notes.txt", + "some text", + ); + expect(result).toBeNull(); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("returns null for empty content", async () => { + const { deps, parseSpy } = makeDeps({}); + const result = await enrichFileForAgent(deps, "/tmp/code.ts", ""); + expect(result).toBeNull(); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("returns null for content > 1MB", async () => { + const { deps, parseSpy } = makeDeps({}); + const huge = "x".repeat(1_000_001); + const result = await enrichFileForAgent(deps, "/tmp/code.ts", huge); + expect(result).toBeNull(); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("returns null when language not supported by enricher", async () => { + const { deps, parseSpy } = makeDeps({ isSupported: false }); + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "posthog.capture('x');", + ); + expect(result).toBeNull(); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("returns null when no PostHog calls detected", async () => { + const { deps, enrichFromApiSpy } = makeDeps({ + callsCount: 0, + initCallsCount: 0, + }); + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "posthog.capture('x');", + ); + expect(result).toBeNull(); + expect(enrichFromApiSpy).not.toHaveBeenCalled(); + }); + + test("returns null and skips parse when content has no posthog reference", async () => { + const { deps, parseSpy } = makeDeps({}); + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "const x = 1;\nfunction foo() {}", + ); + expect(result).toBeNull(); + expect(parseSpy).not.toHaveBeenCalled(); + }); + + test("returns null when getApiKey yields empty string", async () => { + const { deps, enrichFromApiSpy } = makeDeps({ getApiKey: () => "" }); + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "posthog.capture('x');", + ); + expect(result).toBeNull(); + expect(enrichFromApiSpy).not.toHaveBeenCalled(); + }); + + test("returns null when toInlineComments produces no change", async () => { + const original = "posthog.capture('x');"; + const { deps } = makeDeps({ toInlineCommentsReturn: original }); + const result = await enrichFileForAgent(deps, "/tmp/code.ts", original); + expect(result).toBeNull(); + }); + + test("returns null and logs debug when enricher throws", async () => { + const logger = { debug: vi.fn() }; + const { deps } = makeDeps({ parseRejects: new Error("boom") }); + deps.logger = logger as unknown as FileEnrichmentDeps["logger"]; + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "posthog.capture('x');", + ); + expect(result).toBeNull(); + expect(logger.debug).toHaveBeenCalledWith( + "File enrichment failed", + expect.objectContaining({ filePath: "/tmp/code.ts" }), + ); + }); + + test("returns enriched string when happy path completes", async () => { + const { deps, enrichFromApiSpy } = makeDeps({ + toInlineCommentsReturn: "posthog.capture('x'); // [PostHog] Event: \"x\"", + }); + const result = await enrichFileForAgent( + deps, + "/tmp/code.ts", + "posthog.capture('x');", + ); + expect(result).toBe("posthog.capture('x'); // [PostHog] Event: \"x\""); + expect(enrichFromApiSpy).toHaveBeenCalledWith( + expect.objectContaining({ + apiKey: "phx_test", + host: "https://test.posthog.com", + projectId: 1, + }), + ); + }); +}); diff --git a/packages/agent/src/enrichment/file-enricher.ts b/packages/agent/src/enrichment/file-enricher.ts new file mode 100644 index 000000000..5b9eae26d --- /dev/null +++ b/packages/agent/src/enrichment/file-enricher.ts @@ -0,0 +1,82 @@ +import * as path from "node:path"; +import { EXT_TO_LANG_ID, PostHogEnricher } from "@posthog/enricher"; +import type { PostHogAPIConfig } from "../types"; +import type { Logger } from "../utils/logger"; + +export interface FileEnrichmentDeps { + enricher: PostHogEnricher; + apiConfig: PostHogAPIConfig; + logger?: Logger; +} + +export interface Enrichment { + deps: FileEnrichmentDeps; + dispose(): void; +} + +export function createEnrichment( + apiConfig: PostHogAPIConfig | undefined, + logger?: Logger, +): Enrichment | undefined { + if (!apiConfig) return undefined; + const enricher = new PostHogEnricher(); + return { + deps: { enricher, apiConfig, logger }, + dispose: () => enricher.dispose(), + }; +} + +const MAX_ENRICHMENT_BYTES = 1_000_000; + +export async function enrichFileForAgent( + deps: FileEnrichmentDeps, + filePath: string, + content: string, +): Promise { + 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; + + try { + const parsed = await deps.enricher.parse(content, langId); + if (parsed.calls.length === 0 && parsed.initCalls.length === 0) { + return null; + } + + const apiKey = await deps.apiConfig.getApiKey(); + if (!apiKey) return null; + + const enriched = await parsed.enrichFromApi({ + apiKey, + host: deps.apiConfig.apiUrl, + projectId: deps.apiConfig.projectId, + timeoutMs: 5_000, + }); + + const annotated = enriched.toInlineComments(); + if (annotated === content) { + deps.logger?.debug("File enrichment produced no changes", { + filePath, + calls: parsed.calls.length, + }); + return null; + } + deps.logger?.debug("File enriched", { + filePath, + calls: parsed.calls.length, + }); + return annotated; + } catch (err) { + const detail = + err instanceof Error + ? { message: err.message, name: err.name, stack: err.stack } + : { value: String(err) }; + deps.logger?.debug("File enrichment failed", { filePath, ...detail }); + return null; + } +} diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 3399463f6..d9cddeb21 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -152,6 +152,12 @@ export interface AgentConfig { skipLogPersistence?: boolean; /** Local cache path for instant log loading (e.g., ~/.posthog-code) */ localCachePath?: string; + /** + * Annotate files the agent reads with PostHog enrichment (event volume, + * flag rollout/staleness, experiment links). Defaults to enabled when + * `posthog` config is present; set `{ enabled: false }` to opt out. + */ + enricher?: { enabled?: boolean }; debug?: boolean; onLog?: OnLogCallback; } diff --git a/packages/enricher/src/comment-formatter.ts b/packages/enricher/src/comment-formatter.ts index b05350186..0bffa67ee 100644 --- a/packages/enricher/src/comment-formatter.ts +++ b/packages/enricher/src/comment-formatter.ts @@ -44,6 +44,28 @@ function formatEventComment(event: EnrichedEvent): string { return parts.join(" \u2014 "); } +function buildCommentBody( + item: EnrichedListItem, + enrichedFlags: Map, + enrichedEvents: Map, +): string | null { + if (item.type === "flag") { + const flag = enrichedFlags.get(item.name); + if (flag) return formatFlagComment(flag); + return null; + } + if (item.type === "event") { + const event = enrichedEvents.get(item.name); + if (event) return formatEventComment(event); + if (item.detail) return `Event: ${item.detail}`; + return null; + } + if (item.type === "init") { + return `Init: token "${item.name}"`; + } + return null; +} + export function formatComments( source: string, languageId: string, @@ -59,30 +81,41 @@ export function formatComments( for (const item of sorted) { const targetLine = item.line + offset; + const body = buildCommentBody(item, enrichedFlags, enrichedEvents); + if (!body) continue; - let comment: string | null = null; + const comment = `${prefix} [PostHog] ${body}`; + const indent = lines[targetLine]?.match(/^(\s*)/)?.[1] ?? ""; + lines.splice(targetLine, 0, `${indent}${comment}`); + offset++; + } - if (item.type === "flag") { - const flag = enrichedFlags.get(item.name); - if (flag) { - comment = `${prefix} [PostHog] ${formatFlagComment(flag)}`; - } - } else if (item.type === "event") { - const event = enrichedEvents.get(item.name); - if (event) { - comment = `${prefix} [PostHog] ${formatEventComment(event)}`; - } else if (item.detail) { - comment = `${prefix} [PostHog] Event: ${item.detail}`; - } - } else if (item.type === "init") { - comment = `${prefix} [PostHog] Init: token "${item.name}"`; - } + return lines.join("\n"); +} - if (comment) { - const indent = lines[targetLine]?.match(/^(\s*)/)?.[1] ?? ""; - lines.splice(targetLine, 0, `${indent}${comment}`); - offset++; - } +export function formatInlineComments( + source: string, + languageId: string, + items: EnrichedListItem[], + enrichedFlags: Map, + enrichedEvents: Map, +): string { + const prefix = commentPrefix(languageId); + const lines = source.split("\n"); + const byLine = new Map(); + + for (const item of items) { + const body = buildCommentBody(item, enrichedFlags, enrichedEvents); + if (!body) continue; + const arr = byLine.get(item.line) ?? []; + arr.push(body); + byLine.set(item.line, arr); + } + + for (const [lineIdx, bodies] of byLine) { + if (lineIdx < 0 || lineIdx >= lines.length) continue; + const suffix = ` ${prefix} [PostHog] ${bodies.join(" | ")}`; + lines[lineIdx] = `${lines[lineIdx]}${suffix}`; } return lines.join("\n"); diff --git a/packages/enricher/src/enriched-result.ts b/packages/enricher/src/enriched-result.ts index 33d776c4a..7a75193e2 100644 --- a/packages/enricher/src/enriched-result.ts +++ b/packages/enricher/src/enriched-result.ts @@ -1,4 +1,4 @@ -import { formatComments } from "./comment-formatter.js"; +import { formatComments, formatInlineComments } from "./comment-formatter.js"; import { classifyFlagType, extractRollout, @@ -162,4 +162,24 @@ export class EnrichedResult { eventLookup, ); } + + toInlineComments(): string { + const flagLookup = new Map(); + for (const f of this.flags) { + flagLookup.set(f.flagKey, f); + } + + const eventLookup = new Map(); + for (const e of this.events) { + eventLookup.set(e.eventName, e); + } + + return formatInlineComments( + this.parsed.source, + this.parsed.languageId, + this.toList(), + flagLookup, + eventLookup, + ); + } } diff --git a/packages/enricher/src/enricher.test.ts b/packages/enricher/src/enricher.test.ts index a9249790a..af656d1f0 100644 --- a/packages/enricher/src/enricher.test.ts +++ b/packages/enricher/src/enricher.test.ts @@ -282,6 +282,64 @@ describeWithGrammars("PostHogEnricher", () => { expect(annotated).toContain("# [PostHog]"); }); + test("toInlineComments appends to the same line and preserves line count", async () => { + const code = [ + `posthog.capture('purchase');`, + `posthog.getFeatureFlag('my-flag');`, + ].join("\n"); + + const result = await enricher.parse(code, "javascript"); + mockApiResponses({ + flags: [makeFlag("my-flag")], + eventDefs: [makeEventDef("purchase", { verified: true })], + }); + const enriched = await result.enrichFromApi(API_CONFIG); + + const annotated = enriched.toInlineComments(); + const lines = annotated.split("\n"); + + expect(lines).toHaveLength(2); + expect(lines[0]).toMatch(/^posthog\.capture\('purchase'\);.*\[PostHog\]/); + expect(lines[0]).toContain(`Event: "purchase"`); + expect(lines[1]).toMatch( + /^posthog\.getFeatureFlag\('my-flag'\);.*\[PostHog\]/, + ); + expect(lines[1]).toContain(`Flag: "my-flag"`); + }); + + test("toInlineComments uses # for Python", async () => { + const code = `posthog.get_feature_flag('my-flag')`; + const result = await enricher.parse(code, "python"); + + mockApiResponses({ flags: [makeFlag("my-flag")] }); + const enriched = await result.enrichFromApi(API_CONFIG); + + const annotated = enriched.toInlineComments(); + expect(annotated).toContain("# [PostHog]"); + expect(annotated.split("\n")).toHaveLength(1); + }); + + test("toInlineComments combines multiple calls on the same line", async () => { + const code = `posthog.capture('a'); posthog.capture('b');`; + const result = await enricher.parse(code, "javascript"); + + mockApiResponses({ + eventDefs: [ + makeEventDef("a", { verified: true }), + makeEventDef("b", { verified: true }), + ], + }); + const enriched = await result.enrichFromApi(API_CONFIG); + + const annotated = enriched.toInlineComments(); + const lines = annotated.split("\n"); + + expect(lines).toHaveLength(1); + expect(lines[0]).toContain(`Event: "a"`); + expect(lines[0]).toContain(`Event: "b"`); + expect(lines[0]).toContain(" | "); + }); + test("enrichedEvents surfaces stats, lastSeenAt, and tags", async () => { const code = `posthog.capture('purchase');`; const result = await enricher.parse(code, "javascript"); @@ -431,7 +489,7 @@ describeWithGrammars("PostHogEnricher", () => { vi.unstubAllGlobals(); }); - test("rejects on 401 unauthorized", async () => { + test("tolerates 401 unauthorized by returning empty enrichment", async () => { const code = `posthog.getFeatureFlag('my-flag');`; const result = await enricher.parse(code, "javascript"); @@ -440,12 +498,11 @@ describeWithGrammars("PostHogEnricher", () => { vi.fn(async () => new Response("Unauthorized", { status: 401 })), ); - await expect(result.enrichFromApi(API_CONFIG)).rejects.toThrow( - /PostHog API error: 401/, - ); + const enriched = await result.enrichFromApi(API_CONFIG); + expect(enriched.flags[0].flag).toBeUndefined(); }); - test("rejects on 500 server error", async () => { + test("tolerates 500 server error by returning empty enrichment", async () => { const code = `posthog.getFeatureFlag('my-flag');`; const result = await enricher.parse(code, "javascript"); @@ -456,12 +513,11 @@ describeWithGrammars("PostHogEnricher", () => { ), ); - await expect(result.enrichFromApi(API_CONFIG)).rejects.toThrow( - /PostHog API error: 500/, - ); + const enriched = await result.enrichFromApi(API_CONFIG); + expect(enriched.flags[0].flag).toBeUndefined(); }); - test("rejects on network failure", async () => { + test("tolerates network failure by returning empty enrichment", async () => { const code = `posthog.getFeatureFlag('my-flag');`; const result = await enricher.parse(code, "javascript"); @@ -472,12 +528,11 @@ describeWithGrammars("PostHogEnricher", () => { }), ); - await expect(result.enrichFromApi(API_CONFIG)).rejects.toThrow( - "fetch failed", - ); + const enriched = await result.enrichFromApi(API_CONFIG); + expect(enriched.flags[0].flag).toBeUndefined(); }); - test("rejects on malformed JSON response", async () => { + test("tolerates malformed JSON response by returning empty enrichment", async () => { const code = `posthog.getFeatureFlag('my-flag');`; const result = await enricher.parse(code, "javascript"); @@ -492,7 +547,8 @@ describeWithGrammars("PostHogEnricher", () => { ), ); - await expect(result.enrichFromApi(API_CONFIG)).rejects.toThrow(); + const enriched = await result.enrichFromApi(API_CONFIG); + expect(enriched.flags[0].flag).toBeUndefined(); }); }); }); diff --git a/packages/enricher/src/parse-result.ts b/packages/enricher/src/parse-result.ts index 858d41031..5a14f396b 100644 --- a/packages/enricher/src/parse-result.ts +++ b/packages/enricher/src/parse-result.ts @@ -1,8 +1,10 @@ import { EnrichedResult } from "./enriched-result.js"; +import { warn } from "./log.js"; import { PostHogApi } from "./posthog-api.js"; import type { CapturedEvent, EnricherApiConfig, + EventStats, FlagAssignment, FlagCheck, FunctionInfo, @@ -102,17 +104,42 @@ export class ParseResult { const flagKeys = this.flagKeys; const eventNames = this.eventNames; - const [allFlags, allExperiments, allEventDefs, eventStats] = - await Promise.all([ - flagKeys.length > 0 ? api.getFeatureFlags() : Promise.resolve([]), - flagKeys.length > 0 ? api.getExperiments() : Promise.resolve([]), - eventNames.length > 0 - ? api.getEventDefinitions(eventNames) - : Promise.resolve([]), - eventNames.length > 0 - ? api.getEventStats(eventNames) - : Promise.resolve(new Map()), - ]); + const settled = await Promise.allSettled([ + flagKeys.length > 0 ? api.getFeatureFlags() : Promise.resolve([]), + flagKeys.length > 0 ? api.getExperiments() : Promise.resolve([]), + eventNames.length > 0 + ? api.getEventDefinitions(eventNames) + : Promise.resolve([]), + eventNames.length > 0 + ? api.getEventStats(eventNames) + : Promise.resolve(new Map()), + ]); + + const [flagsResult, experimentsResult, eventDefsResult, eventStatsResult] = + settled; + + const labels = [ + "getFeatureFlags", + "getExperiments", + "getEventDefinitions", + "getEventStats", + ]; + settled.forEach((r, i) => { + if (r.status === "rejected") { + warn(`enricher: ${labels[i]} failed`, r.reason); + } + }); + + const allFlags = + flagsResult.status === "fulfilled" ? flagsResult.value : []; + const allExperiments = + experimentsResult.status === "fulfilled" ? experimentsResult.value : []; + const allEventDefs = + eventDefsResult.status === "fulfilled" ? eventDefsResult.value : []; + const eventStats = + eventStatsResult.status === "fulfilled" + ? eventStatsResult.value + : new Map(); const flagKeySet = new Set(flagKeys); const flags = new Map( diff --git a/packages/enricher/src/parser-manager.ts b/packages/enricher/src/parser-manager.ts index f26a4bd0f..ac9afa637 100644 --- a/packages/enricher/src/parser-manager.ts +++ b/packages/enricher/src/parser-manager.ts @@ -1,3 +1,4 @@ +import { existsSync } from "node:fs"; import * as path from "node:path"; import { fileURLToPath } from "node:url"; import Parser from "web-tree-sitter"; @@ -8,9 +9,14 @@ import type { DetectionConfig } from "./types.js"; import { DEFAULT_CONFIG } from "./types.js"; function resolveGrammarsDir(): string { - // Works from both dist/ (built) and src/ (tests) — both are one level below package root const thisFile = fileURLToPath(import.meta.url); - return path.join(path.dirname(thisFile), "..", "grammars"); + const dir = path.dirname(thisFile); + const candidates = [ + path.join(dir, "..", "grammars"), + path.join(dir, "grammars"), + path.join(dir, "..", "..", "grammars"), + ]; + return candidates.find((p) => existsSync(p)) ?? candidates[0]; } export class ParserManager { diff --git a/packages/enricher/src/posthog-api.ts b/packages/enricher/src/posthog-api.ts index b81b6704b..d34b38ab6 100644 --- a/packages/enricher/src/posthog-api.ts +++ b/packages/enricher/src/posthog-api.ts @@ -87,6 +87,9 @@ export class PostHogApi { return new Map(); } + // HogQL over `/query/` rejects typed placeholders (`{name:Type}`) and + // placeholder values in INTERVAL, so `days` is inlined (clamped). + const days = Math.max(1, Math.min(365, Math.floor(daysBack))); const query = ` SELECT event, @@ -94,8 +97,8 @@ export class PostHogApi { count(DISTINCT person_id) AS unique_users, max(timestamp) AS last_seen FROM events - WHERE event IN ({eventNames:Array(String)}) - AND timestamp >= now() - INTERVAL {daysBack:Int32} DAY + WHERE event IN {eventNames} + AND timestamp >= now() - INTERVAL ${days} DAY GROUP BY event `; @@ -105,7 +108,7 @@ export class PostHogApi { query: { kind: "HogQLQuery", query, - values: { eventNames, daysBack }, + values: { eventNames }, }, }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cc3e99d73..ba19fc862 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -657,6 +657,9 @@ importers: '@opentelemetry/semantic-conventions': specifier: ^1.28.0 version: 1.39.0 + '@posthog/enricher': + specifier: workspace:* + version: link:../enricher '@types/jsonwebtoken': specifier: ^9.0.10 version: 9.0.10