refactor(heartbeat): extract run retrieval and session state (#15601)

## Thinking Path

> - Paperclip is the open source app people use to manage AI agents for
work.
> - Heartbeat runs retain task sessions, usage totals, and bounded run
data.
> - After the preparation extraction, the main service still has over
24,000 lines.
> - Run reads and session operations share one database and have no
dispatch dependencies.
> - This pull request moves that group into one module and preserves its
callers.
> - The benefit is a smaller run engine and one place to maintain saved
run and session state.

## Linked Issues or Issue Description

**What existing behavior does this improve?**

It improves the structure and test coverage of heartbeat run retrieval
and session state.

**Current behavior**

`server/src/services/heartbeat.ts` has 24,033 lines on the merged base.
It mixes run projections, session queries, resume rules, compaction, and
usage helpers with execution and recovery.

**Proposed behavior**

Move that group into `server/src/services/heartbeat/run-state.ts`. The
main file loses 1,770 lines and ends at 22,263 lines. Keep current
behavior, public exports, and service methods. This continues #15568,
#15573, #15578, and #15591.

**Reason and benefit**

One domain extraction makes useful progress toward a few manageable
files. Run reads and saved session state can be reviewed together
without searching the run engine. A search found no duplicate open
extraction. Related session changes include #15142, #14661, and #14670;
this PR does not apply their proposed behavior changes.

**Breaking changes**

None. Existing public helpers retain their identity. Service method
names and results stay the same.

## What Changed

- Move bounded run projections, encoding checks, and run/event reads
into `heartbeat/run-state.ts`.
- Move task session persistence, explicit resumes, reset rules,
compaction, and usage/billing helpers into the same module.
- Bind database operations through `createHeartbeatRunState(db)`.
Construction does no database work. The encoding cache stays local to
each service instance.
- Keep execution order, status transitions, session-goal recovery, cost
accounting writes, dispatch, and cancellation in `heartbeat.ts`.
- Preserve all 140 exports. A syntax-tree comparison confirms that 43
moved function bodies, 22 declarations, and nine service method bodies
are unchanged. The copied private string helper is unchanged.
Surrounding orchestration is unchanged apart from the factory binding.
- Add eight tests for legacy export identity, database binding, encoding
cache isolation, scoped resumes, cumulative usage, compaction, and stale
or cancelled conversation session writes.
- Document the module boundary in `doc/DEVELOPING.md`.

## Verification

- Before and after extraction, eight existing suites pass: 235 tests.
They cover workspace/session helpers, runtime state, cost accounting,
ledger attribution, task reset, timer reset, run lists, and run privacy.
- The new module suite passes: eight tests, including five real
PostgreSQL cases. Total focused coverage: 243 tests across nine suites.
- Commands: `pnpm exec vitest run
server/src/services/heartbeat/run-state.test.ts` and `pnpm exec vitest
run server/src/__tests__/heartbeat-runtime-state.test.ts
server/src/__tests__/heartbeat-cost-accounting.test.ts
server/src/__tests__/heartbeat-ledger-billing-code.test.ts
server/src/__tests__/heartbeat-task-session-reset.test.ts
server/src/__tests__/heartbeat-workspace-session.test.ts
server/src/__tests__/heartbeat-timer-wake-session-reset-pf4.test.ts
server/src/__tests__/heartbeat-list.test.ts
server/src/__tests__/heartbeat-run-privacy-routes.test.ts`.
- `pnpm --filter @paperclipai/server typecheck` passes.
- `pnpm -r typecheck` and `pnpm build` pass.
- The full local `pnpm test:run` did not complete and was stopped with
SIGINT after it reported an OAuth scope health-check failure and a
workspace-runtime timeout. All six OAuth variants and the workspace case
pass in isolation. A single rerun of the 413-case tool-access suite
passed the OAuth case but hit a different Notion callback timeout; that
Notion case also passes in isolation. The full local suite is not
claimed as passing.
- Greptile is 5/5 with no findings or inline comments on head
`2ca0908cfc112dd6ae26e2565cb3611b5bec52ba`. Its exact-head check
completed successfully.
- CI remains blocked on hosted-runner assignment after more than ten
minutes. [`ci / Select trusted
runner`](https://github.com/paperclipai/paperclip/actions/runs/37828454507/job/113487287178)
is queued for `ubuntu-latest` with no runner assigned. All seven
completed review/security checks pass; two optional Storybook checks are
skipped. The full CI matrix has not started. It must complete before
merge.

## Risks

- Wiring errors could bind reads to the wrong database or share the
encoding cache. Direct factory tests cover database and cache isolation.
Real PostgreSQL suites cover query and transaction behavior.
- Session writes must retain their issue row locks and generation
fences. The original function bodies are unchanged. New database tests
reject stale and cancelled conversation writes and clears.
- Existing SQL_ASCII safeguards and bounded projections remain in place.
No schema, API contract, migration, lockfile, or workflow changes are
included.

## Model Used

OpenAI GPT-6 via Codex. The exact serving model ID and context window
were not exposed in this session. Used reasoning, repository inspection,
shell tools, and code execution.

## Checklist

- [x] I have included a thinking path that traces from project context
to this change
- [x] I have specified the model used (with version and capability
details)
- [x] I have checked ROADMAP.md and confirmed this PR does not duplicate
planned core work
- [x] I have searched GitHub for duplicate or related PRs and linked
them above
- [x] I have either (a) linked existing issues with `Fixes: #` / `Closes
#` / `Refs #` OR (b) described the issue in-PR following the relevant
issue template
- [x] I have not referenced internal/instance-local Paperclip issues or
links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip`
URLs)
- [x] My branch name describes the change (e.g. `docs/...`, `fix/...`)
and contains no internal Paperclip ticket id or instance-derived details
- [x] I have run tests locally and they pass
- [x] I have added or updated tests where applicable
- [x] I have updated relevant documentation to reflect my changes
- [x] I have considered and documented any risks above
- [ ] All Paperclip CI gates are green
- [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups
- [x] I will address all Greptile and reviewer comments before
requesting merge

Co-authored-by: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip authored and GitHub committed 2026-10-08 14:53:48 -05:00
1 parent 8e59efc50b
commit 4fe45bf8fe
4 files changed
+2228 -1844

No files matched your search

+9
View File
@@ -1245,6 +1245,15 @@ dispatch, retries, cancellation, and execution order, and re-exports the existin
public helpers and configuration-incomplete error class. Keep preparation policy
changes in this module and its tests.
Run retrieval and session state are in `server/src/services/heartbeat/run-state.ts`.
It owns bounded run projections, database encoding checks, task session reads and
writes, explicit resumes, session compaction, and usage/billing helpers.
`createHeartbeatRunState(db)` binds these operations without doing database work
during construction. The encoding-check cache belongs to each factory instance.
`heartbeat.ts` keeps run execution, session-goal recovery, cost accounting writes,
and status transitions, and re-exports the existing public helpers. Keep session
policy changes separate from run orchestration changes.
## Wake Context Delivery
Built-in adapters deliver wake context through the run prompt, including structured
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,222 @@
import { randomUUID } from "node:crypto";
import type { Db } from "@paperclipai/db";
import {
agents,
agentTaskSessions,
companies,
createDb,
heartbeatRuns,
issues,
} from "@paperclipai/db";
import { eq, type SQL } from "drizzle-orm";
import { PgDialect } from "drizzle-orm/pg-core";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import {
getEmbeddedPostgresTestSupport,
startEmbeddedPostgresTestDatabase,
} from "../../__tests__/helpers/embedded-postgres.js";
import * as legacy from "../heartbeat.js";
import * as extracted from "./run-state.js";
// SQL_ASCII needs a separate database cluster. Inspect its selected SQL here;
// the PostgreSQL tests below exercise the session queries and transactions.
function readingDatabase(serverEncoding: string, ...results: unknown[][]) {
const scopes: unknown[][] = [];
const columns: Record<string, unknown>[] = [];
const execute = vi.fn(async () => [{ server_encoding: serverEncoding }]);
const select = vi.fn((projection: Record<string, unknown> = {}) => {
columns.push(projection);
const rows = results.shift() ?? [];
const query = {
from: () => query,
where: (condition: SQL) => {
scopes.push(new PgDialect().sqlToQuery(condition).params);
return query;
},
orderBy: () => query,
limit: () => query,
then: (resolve: (rows: unknown[]) => unknown) => Promise.resolve(rows).then(resolve),
};
return query;
});
return { db: { execute, select } as unknown as Db, execute, select, columns, scopes };
}
describe("heartbeat run state module", () => {
it("preserves the legacy public helper identities", () => {
expect(legacy.resolveNextSessionState).toBe(extracted.resolveNextSessionState);
expect(legacy.buildExplicitResumeSessionOverride).toBe(extracted.buildExplicitResumeSessionOverride);
expect(legacy.normalizeAdapterRunUsage).toBe(extracted.normalizeAdapterRunUsage);
expect(legacy.resolveLedgerScopeForRun).toBe(extracted.resolveLedgerScopeForRun);
expect(legacy.summarizeHeartbeatRunListResultJson).toBe(extracted.summarizeHeartbeatRunListResultJson);
});
it("constructs without queries and binds session lookups to each database", async () => {
const first = readingDatabase("UTF8", [{ sessionDisplayId: "first" }]);
const second = readingDatabase("UTF8", [{ sessionDisplayId: "second" }]);
const firstState = extracted.createHeartbeatRunState(first.db);
const secondState = extracted.createHeartbeatRunState(second.db);
expect(first.select).not.toHaveBeenCalled();
expect(first.execute).not.toHaveBeenCalled();
expect(second.select).not.toHaveBeenCalled();
expect(second.execute).not.toHaveBeenCalled();
const [one, two] = await Promise.all([
firstState.getTaskSession("company-one", "agent-one", "codex_local", "task-one"),
secondState.getTaskSession("company-two", "agent-two", "claude_local", "task-two"),
]);
expect(one?.sessionDisplayId).toBe("first");
expect(two?.sessionDisplayId).toBe("second");
expect(first.scopes[0]).toEqual(["company-one", "agent-one", "codex_local", "task-one"]);
expect(second.scopes[0]).toEqual(["company-two", "agent-two", "claude_local", "task-two"]);
});
it("keeps encoding caches local to each factory and preserves full-result reads", async () => {
const ascii = readingDatabase("SQL_ASCII");
const utf8 = readingDatabase("UTF8");
const asciiState = extracted.createHeartbeatRunState(ascii.db);
const utf8State = extracted.createHeartbeatRunState(utf8.db);
await Promise.all([asciiState.getRun("ascii-one"), asciiState.getRun("ascii-two"), utf8State.getRun("utf8")]);
const projectedSql = (column: unknown) => new PgDialect().sqlToQuery((column as SQL.Aliased).sql).sql;
expect(ascii.execute).toHaveBeenCalledTimes(1);
expect(utf8.execute).toHaveBeenCalledTimes(1);
expect(projectedSql(ascii.columns[0].resultJson)).toBe("NULL");
expect(projectedSql(utf8.columns[0].resultJson)).toContain("pg_column_size");
await asciiState.getRun("evidence", { includeExecutionEvidence: true });
expect(projectedSql(ascii.columns[2].resultJson)).toContain("startupPreparationSettledAt");
const raw = readingDatabase("SQL_ASCII");
await extracted.createHeartbeatRunState(raw.db).getRun("raw", { unsafeFullResultJson: true });
expect(raw.execute).not.toHaveBeenCalled();
expect(raw.columns[0].resultJson).toBe(heartbeatRuns.resultJson);
});
});
const postgresSupport = await getEmbeddedPostgresTestSupport();
const describePostgres = postgresSupport.supported ? describe : describe.skip;
if (!postgresSupport.supported) {
console.warn(`Skipping run-state PostgreSQL tests: ${postgresSupport.reason ?? "unsupported host"}`);
}
describePostgres("heartbeat run state database wiring", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("heartbeat-run-state-");
db = createDb(database.connectionString);
}, 20_000);
afterAll(async () => {
await database?.cleanup();
});
async function fixture() {
const companyId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Run state", issuePrefix: `S${companyId.slice(0, 8)}` });
const [agent] = await db.insert(agents).values({
companyId, name: "Session agent", adapterType: "codex_local",
}).returning();
const [issue] = await db.insert(issues).values({ companyId, title: "Session task", assigneeAgentId: agent.id }).returning();
return { companyId, agent, issue, state: extracted.createHeartbeatRunState(db) };
}
it("resumes only the selected agent's run and task session", async () => {
const { companyId, agent, issue, state } = await fixture();
const foreign = await fixture();
const [run] = await db.insert(heartbeatRuns).values({
companyId, agentId: agent.id, invocationSource: "assignment", status: "failed",
contextSnapshot: { issueId: issue.id }, sessionIdAfter: "our-session",
}).returning();
const [foreignRun] = await db.insert(heartbeatRuns).values({
companyId: foreign.companyId, agentId: foreign.agent.id, invocationSource: "assignment",
contextSnapshot: { issueId: foreign.issue.id }, sessionIdAfter: "foreign-session",
}).returning();
await state.upsertTaskSession({
companyId, agentId: agent.id, adapterType: agent.adapterType, taskKey: issue.id,
sessionParamsJson: { sessionId: "our-session", cwd: "/saved-workspace" },
sessionDisplayId: "our-session", lastRunId: run.id, lastError: null,
});
expect(await state.resolveExplicitResumeSessionOverride(agent, { resumeFromRunId: foreignRun.id }, issue.id)).toBeNull();
expect(await state.resolveExplicitResumeSessionOverride(agent, { resumeFromRunId: run.id }, issue.id)).toMatchObject({
resumeFromRunId: run.id, issueId: issue.id, taskKey: issue.id,
sessionDisplayId: "our-session", sessionParams: { sessionId: "our-session", cwd: "/saved-workspace" },
});
expect(await state.resolveSessionBeforeForWakeup(agent, issue.id)).toBe("our-session");
expect(await state.getRuntimeStateWithSessions(agent.id)).toMatchObject({
companyId, sessionDisplayId: "our-session", sessionParamsJson: { sessionId: "our-session" },
});
expect((await state.listRuns(companyId)).map((row) => row.id)).toEqual([run.id]);
});
it("subtracts prior cumulative usage while excluding the current run", async () => {
const { companyId, agent, state } = await fixture();
await db.insert(heartbeatRuns).values({
companyId, agentId: agent.id, invocationSource: "assignment", sessionIdAfter: "cumulative",
createdAt: new Date("2026-01-01T00:00:00Z"),
usageJson: { rawInputTokens: 800, rawCachedInputTokens: 100, rawOutputTokens: 80 },
});
const [current] = await db.insert(heartbeatRuns).values({
companyId, agentId: agent.id, invocationSource: "assignment", sessionIdAfter: "cumulative",
createdAt: new Date("2026-01-02T00:00:00Z"),
usageJson: { rawInputTokens: 1000, rawCachedInputTokens: 125, rawOutputTokens: 100 },
}).returning();
const rawUsage = { inputTokens: 1000, cachedInputTokens: 125, outputTokens: 100 };
const input = { agentId: agent.id, runId: current.id, sessionId: "cumulative", rawUsage };
expect(await state.resolveNormalizedUsageForSession({ ...input, usageBasis: "session_cumulative" })).toEqual({
normalizedUsage: { inputTokens: 200, cachedInputTokens: 25, outputTokens: 20 },
previousRawUsage: { inputTokens: 800, cachedInputTokens: 100, outputTokens: 80 },
derivedFromSessionTotals: true,
});
expect(await state.resolveNormalizedUsageForSession({ ...input, usageBasis: "per_run" })).toEqual({
normalizedUsage: rawUsage, previousRawUsage: null, derivedFromSessionTotals: false,
});
});
it("uses the selected session's usage and summary for compaction", async () => {
const { companyId, agent, issue, state } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({
companyId, agentId: agent.id, invocationSource: "assignment", sessionIdAfter: "compact",
usageJson: { rawInputTokens: 90, rawCachedInputTokens: 20, rawInputIncludesCached: false },
resultJson: { summary: "Continue our task", providerPrivateData: "omit this" },
}).returning();
const foreign = await fixture();
await db.insert(heartbeatRuns).values({
companyId: foreign.companyId, agentId: foreign.agent.id, invocationSource: "assignment",
sessionIdAfter: "compact", resultJson: { summary: "Foreign task" },
});
const decision = await state.evaluateSessionCompaction({
agent: { ...agent, runtimeConfig: { heartbeat: { sessionCompaction: { maxRawInputTokens: 100 } } } },
sessionId: "compact", issueId: issue.id,
});
expect(decision).toMatchObject({ rotate: true, reason: "session raw input reached 110 tokens (threshold 100)", previousRunId: run.id });
expect(decision.handoffMarkdown).toContain("Continue our task");
expect(decision.handoffMarkdown).not.toContain("Foreign task");
expect(decision.handoffMarkdown).not.toContain("omit this");
});
it.each(["superseded", "cancelled"] as const)("fences %s conversation session writes and clears", async (scenario) => {
const { companyId, agent, issue, state } = await fixture();
await db.update(issues).set({
conversationAgentId: agent.id, conversationUserId: "test-user", conversationState: "active",
conversationSessionGeneration: 2,
}).where(eq(issues.id, issue.id));
const [run] = await db.insert(heartbeatRuns).values({
companyId, agentId: agent.id, invocationSource: "assignment", status: "succeeded",
contextSnapshot: { conversationSessionGeneration: 2 },
}).returning();
const input = {
companyId, agentId: agent.id, adapterType: agent.adapterType, taskKey: issue.id,
sessionParamsJson: { sessionId: "retained" }, sessionDisplayId: "retained", lastRunId: run.id, lastError: null,
};
const saved = await state.upsertTaskSession(input);
expect(saved?.sessionDisplayId).toBe("retained");
if (scenario === "superseded") {
await db.update(issues).set({ conversationSessionGeneration: 3 }).where(eq(issues.id, issue.id));
} else {
await db.update(heartbeatRuns).set({ status: "cancelled" }).where(eq(heartbeatRuns.id, run.id));
}
expect(await state.upsertTaskSession({ ...input, sessionDisplayId: "stale" })).toBeNull();
expect(await state.clearTaskSessions(companyId, agent.id, { taskKey: issue.id, expectedRunId: run.id })).toBe(0);
const [retained] = await db.select().from(agentTaskSessions).where(eq(agentTaskSessions.id, saved!.id));
expect(retained.sessionDisplayId).toBe("retained");
});
});
File diff suppressed because it is too large. Load diff