fix: use agent role for first heartbeat telemetry

Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
dotta 2026-04-02 09:18:26 -05:00
parent daea94a2ed
commit 85e6371cb6
2 changed files with 36 additions and 3 deletions

View file

@ -1,7 +1,7 @@
import { randomUUID } from "node:crypto"; import { randomUUID } from "node:crypto";
import { spawn, type ChildProcess } from "node:child_process"; import { spawn, type ChildProcess } from "node:child_process";
import { eq } from "drizzle-orm"; import { eq } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import { import {
agents, agents,
agentWakeupRequests, agentWakeupRequests,
@ -16,6 +16,23 @@ import {
startEmbeddedPostgresTestDatabase, startEmbeddedPostgresTestDatabase,
} from "./helpers/embedded-postgres.js"; } from "./helpers/embedded-postgres.js";
import { runningProcesses } from "../adapters/index.ts"; import { runningProcesses } from "../adapters/index.ts";
const mockTelemetryClient = vi.hoisted(() => ({ track: vi.fn() }));
const mockTrackAgentFirstHeartbeat = vi.hoisted(() => vi.fn());
vi.mock("../telemetry.ts", () => ({
getTelemetryClient: () => mockTelemetryClient,
}));
vi.mock("@paperclipai/shared/telemetry", async () => {
const actual = await vi.importActual<typeof import("@paperclipai/shared/telemetry")>(
"@paperclipai/shared/telemetry",
);
return {
...actual,
trackAgentFirstHeartbeat: mockTrackAgentFirstHeartbeat,
};
});
import { heartbeatService } from "../services/heartbeat.ts"; import { heartbeatService } from "../services/heartbeat.ts";
const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip;
@ -43,6 +60,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
}, 20_000); }, 20_000);
afterEach(async () => { afterEach(async () => {
vi.clearAllMocks();
runningProcesses.clear(); runningProcesses.clear();
for (const child of childProcesses) { for (const child of childProcesses) {
child.kill("SIGKILL"); child.kill("SIGKILL");
@ -67,6 +85,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
async function seedRunFixture(input?: { async function seedRunFixture(input?: {
adapterType?: string; adapterType?: string;
agentStatus?: "paused" | "idle" | "running";
runStatus?: "running" | "queued" | "failed"; runStatus?: "running" | "queued" | "failed";
processPid?: number | null; processPid?: number | null;
processLossRetryCount?: number; processLossRetryCount?: number;
@ -94,7 +113,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
companyId, companyId,
name: "CodexCoder", name: "CodexCoder",
role: "engineer", role: "engineer",
status: "paused", status: input?.agentStatus ?? "paused",
adapterType: input?.adapterType ?? "codex_local", adapterType: input?.adapterType ?? "codex_local",
adapterConfig: {}, adapterConfig: {},
runtimeConfig: {}, runtimeConfig: {},
@ -252,4 +271,18 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
expect(run?.errorCode).toBeNull(); expect(run?.errorCode).toBeNull();
expect(run?.error).toBeNull(); expect(run?.error).toBeNull();
}); });
it("tracks the first heartbeat with the agent role instead of adapter type", async () => {
const { runId } = await seedRunFixture({
agentStatus: "running",
includeIssue: false,
});
const heartbeat = heartbeatService(db);
await heartbeat.cancelRun(runId);
expect(mockTrackAgentFirstHeartbeat).toHaveBeenCalledWith(mockTelemetryClient, {
agentRole: "engineer",
});
});
}); });

View file

@ -1832,7 +1832,7 @@ export function heartbeatService(db: Db) {
if (isFirstHeartbeat && updated) { if (isFirstHeartbeat && updated) {
const tc = getTelemetryClient(); const tc = getTelemetryClient();
if (tc) trackAgentFirstHeartbeat(tc, { agentRole: updated.adapterType }); if (tc) trackAgentFirstHeartbeat(tc, { agentRole: updated.role });
} }
if (updated) { if (updated) {