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 { spawn, type ChildProcess } from "node:child_process";
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 {
agents,
agentWakeupRequests,
@ -16,6 +16,23 @@ import {
startEmbeddedPostgresTestDatabase,
} from "./helpers/embedded-postgres.js";
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";
const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip;
@ -43,6 +60,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
}, 20_000);
afterEach(async () => {
vi.clearAllMocks();
runningProcesses.clear();
for (const child of childProcesses) {
child.kill("SIGKILL");
@ -67,6 +85,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
async function seedRunFixture(input?: {
adapterType?: string;
agentStatus?: "paused" | "idle" | "running";
runStatus?: "running" | "queued" | "failed";
processPid?: number | null;
processLossRetryCount?: number;
@ -94,7 +113,7 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
companyId,
name: "CodexCoder",
role: "engineer",
status: "paused",
status: input?.agentStatus ?? "paused",
adapterType: input?.adapterType ?? "codex_local",
adapterConfig: {},
runtimeConfig: {},
@ -252,4 +271,18 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
expect(run?.errorCode).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) {
const tc = getTelemetryClient();
if (tc) trackAgentFirstHeartbeat(tc, { agentRole: updated.adapterType });
if (tc) trackAgentFirstHeartbeat(tc, { agentRole: updated.role });
}
if (updated) {