Merge pull request #1654 from paperclipai/pr/pap-795-agent-runtime

fix(runtime): improve agent recovery and heartbeat operations
This commit is contained in:
Dotta 2026-03-23 19:44:51 -05:00 committed by GitHub
commit f2637e6972
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 1291 additions and 64 deletions

View file

@ -272,6 +272,62 @@ function deriveBundleState(agent: AgentLike): BundleState {
};
}
async function recoverManagedBundleState(agent: AgentLike, state: BundleState): Promise<BundleState> {
const managedRootPath = resolveManagedInstructionsRoot(agent);
const stat = await statIfExists(managedRootPath);
if (!stat?.isDirectory()) return state;
const files = await listFilesRecursive(managedRootPath);
if (files.length === 0) return state;
const recoveredEntryFile = files.includes(state.entryFile)
? state.entryFile
: files.includes(ENTRY_FILE_DEFAULT)
? ENTRY_FILE_DEFAULT
: files[0]!;
if (!state.rootPath) {
return {
...state,
mode: "managed",
rootPath: managedRootPath,
entryFile: recoveredEntryFile,
resolvedEntryPath: path.resolve(managedRootPath, recoveredEntryFile),
};
}
if (state.mode === "external") return state;
const resolvedConfiguredRoot = path.resolve(state.rootPath);
const configuredRootMatchesManaged = resolvedConfiguredRoot === managedRootPath;
const hasEntryMismatch = recoveredEntryFile !== state.entryFile;
if (configuredRootMatchesManaged && !hasEntryMismatch) {
return state;
}
const warnings = [...state.warnings];
if (!configuredRootMatchesManaged) {
warnings.push(
`Recovered managed instructions from disk at ${managedRootPath}; ignoring stale configured root ${state.rootPath}.`,
);
}
if (hasEntryMismatch) {
warnings.push(
`Recovered managed instructions entry file from disk as ${recoveredEntryFile}; previous entry ${state.entryFile} was missing.`,
);
}
return {
...state,
mode: "managed",
rootPath: managedRootPath,
entryFile: recoveredEntryFile,
resolvedEntryPath: path.resolve(managedRootPath, recoveredEntryFile),
warnings,
};
}
function toBundle(agent: AgentLike, state: BundleState, files: AgentInstructionsFileSummary[]): AgentInstructionsBundle {
const nextFiles = [...files];
if (state.legacyPromptTemplateActive && !nextFiles.some((file) => file.path === LEGACY_PROMPT_TEMPLATE_PATH)) {
@ -327,6 +383,36 @@ function applyBundleConfig(
return next;
}
function buildPersistedBundleConfig(
derived: BundleState,
current: BundleState,
options?: { clearLegacyPromptTemplate?: boolean },
): Record<string, unknown> {
const currentRootPath = current.rootPath ? path.resolve(current.rootPath) : null;
const derivedRootPath = derived.rootPath ? path.resolve(derived.rootPath) : null;
const configMatchesRecoveredState =
derived.mode === current.mode
&& derivedRootPath !== null
&& currentRootPath !== null
&& derivedRootPath === currentRootPath
&& derived.entryFile === current.entryFile;
if (configMatchesRecoveredState && !options?.clearLegacyPromptTemplate) {
return current.config;
}
if (!current.rootPath || !current.mode) {
return current.config;
}
return applyBundleConfig(current.config, {
mode: current.mode,
rootPath: current.rootPath,
entryFile: current.entryFile,
clearLegacyPromptTemplate: options?.clearLegacyPromptTemplate,
});
}
async function writeBundleFiles(
rootPath: string,
files: Record<string, string>,
@ -366,7 +452,7 @@ export function syncInstructionsBundleConfigFromFilePath(
export function agentInstructionsService() {
async function getBundle(agent: AgentLike): Promise<AgentInstructionsBundle> {
const state = deriveBundleState(agent);
const state = await recoverManagedBundleState(agent, deriveBundleState(agent));
if (!state.rootPath) return toBundle(agent, state, []);
const stat = await statIfExists(state.rootPath);
if (!stat?.isDirectory()) {
@ -381,7 +467,7 @@ export function agentInstructionsService() {
}
async function readFile(agent: AgentLike, relativePath: string): Promise<AgentInstructionsFileDetail> {
const state = deriveBundleState(agent);
const state = await recoverManagedBundleState(agent, deriveBundleState(agent));
if (relativePath === LEGACY_PROMPT_TEMPLATE_PATH) {
const content = asString(state.config[PROMPT_KEY]);
if (content === null) throw notFound("Instructions file not found");
@ -422,9 +508,14 @@ export function agentInstructionsService() {
agent: AgentLike,
options?: { clearLegacyPromptTemplate?: boolean },
): Promise<{ adapterConfig: Record<string, unknown>; state: BundleState }> {
const current = deriveBundleState(agent);
const derived = deriveBundleState(agent);
const current = await recoverManagedBundleState(agent, derived);
if (current.rootPath && current.mode) {
return { adapterConfig: current.config, state: current };
const adapterConfig = buildPersistedBundleConfig(derived, current, options);
return {
adapterConfig,
state: deriveBundleState({ ...agent, adapterConfig }),
};
}
const managedRoot = resolveManagedInstructionsRoot(agent);
@ -462,7 +553,7 @@ export function agentInstructionsService() {
clearLegacyPromptTemplate?: boolean;
},
): Promise<{ bundle: AgentInstructionsBundle; adapterConfig: Record<string, unknown> }> {
const state = deriveBundleState(agent);
const state = await recoverManagedBundleState(agent, deriveBundleState(agent));
const nextMode = input.mode ?? state.mode ?? "managed";
const nextEntryFile = input.entryFile ? normalizeRelativeFilePath(input.entryFile) : state.entryFile;
let nextRootPath: string;
@ -544,7 +635,8 @@ export function agentInstructionsService() {
bundle: AgentInstructionsBundle;
adapterConfig: Record<string, unknown>;
}> {
const state = deriveBundleState(agent);
const derived = deriveBundleState(agent);
const state = await recoverManagedBundleState(agent, derived);
if (relativePath === LEGACY_PROMPT_TEMPLATE_PATH) {
throw unprocessable("Cannot delete the legacy promptTemplate pseudo-file");
}
@ -555,8 +647,9 @@ export function agentInstructionsService() {
}
const absolutePath = resolvePathWithinRoot(state.rootPath, normalizedPath);
await fs.rm(absolutePath, { force: true });
const bundle = await getBundle(agent);
return { bundle, adapterConfig: state.config };
const adapterConfig = buildPersistedBundleConfig(derived, state);
const bundle = await getBundle({ ...agent, adapterConfig });
return { bundle, adapterConfig };
}
async function exportFiles(agent: AgentLike): Promise<{
@ -564,7 +657,7 @@ export function agentInstructionsService() {
entryFile: string;
warnings: string[];
}> {
const state = deriveBundleState(agent);
const state = await recoverManagedBundleState(agent, deriveBundleState(agent));
if (state.rootPath) {
const stat = await statIfExists(state.rootPath);
if (stat?.isDirectory()) {

View file

@ -132,6 +132,21 @@ export function defaultIssueExecutionWorkspaceSettingsForProject(
};
}
export function issueExecutionWorkspaceModeForPersistedWorkspace(
mode: string | null | undefined,
): IssueExecutionWorkspaceSettings["mode"] {
if (mode === null || mode === undefined) {
return "agent_default";
}
if (mode === "isolated_workspace" || mode === "operator_branch" || mode === "shared_workspace") {
return mode;
}
if (mode === "adapter_managed" || mode === "cloud_sandbox") {
return "agent_default";
}
return "shared_workspace";
}
export function resolveExecutionWorkspaceMode(input: {
projectPolicy: ProjectExecutionWorkspacePolicy | null;
issueSettings: IssueExecutionWorkspaceSettings | null;

View file

@ -45,6 +45,7 @@ import { workspaceOperationService } from "./workspace-operations.js";
import {
buildExecutionWorkspaceAdapterConfig,
gateProjectExecutionWorkspacePolicy,
issueExecutionWorkspaceModeForPersistedWorkspace,
parseIssueExecutionWorkspaceSettings,
parseProjectExecutionWorkspacePolicy,
resolveExecutionWorkspaceMode,
@ -325,6 +326,51 @@ async function resolveLedgerScopeForRun(
};
}
type ResumeSessionRow = {
sessionParamsJson: Record<string, unknown> | null;
sessionDisplayId: string | null;
lastRunId: string | null;
};
export function buildExplicitResumeSessionOverride(input: {
resumeFromRunId: string;
resumeRunSessionIdBefore: string | null;
resumeRunSessionIdAfter: string | null;
taskSession: ResumeSessionRow | null;
sessionCodec: AdapterSessionCodec;
}) {
const desiredDisplayId = truncateDisplayId(
input.resumeRunSessionIdAfter ?? input.resumeRunSessionIdBefore,
);
const taskSessionParams = normalizeSessionParams(
input.sessionCodec.deserialize(input.taskSession?.sessionParamsJson ?? null),
);
const taskSessionDisplayId = truncateDisplayId(
input.taskSession?.sessionDisplayId ??
(input.sessionCodec.getDisplayId ? input.sessionCodec.getDisplayId(taskSessionParams) : null) ??
readNonEmptyString(taskSessionParams?.sessionId),
);
const canReuseTaskSessionParams =
input.taskSession != null &&
(
input.taskSession.lastRunId === input.resumeFromRunId ||
(!!desiredDisplayId && taskSessionDisplayId === desiredDisplayId)
);
const sessionParams =
canReuseTaskSessionParams
? taskSessionParams
: desiredDisplayId
? { sessionId: desiredDisplayId }
: null;
const sessionDisplayId = desiredDisplayId ?? (canReuseTaskSessionParams ? taskSessionDisplayId : null);
if (!sessionDisplayId && !sessionParams) return null;
return {
sessionDisplayId,
sessionParams,
};
}
function normalizeUsageTotals(usage: UsageSummary | null | undefined): UsageTotals | null {
if (!usage) return null;
return {
@ -977,6 +1023,57 @@ export function heartbeatService(db: Db) {
return runtimeForRun?.sessionId ?? null;
}
async function resolveExplicitResumeSessionOverride(
agent: typeof agents.$inferSelect,
payload: Record<string, unknown> | null,
taskKey: string | null,
) {
const resumeFromRunId = readNonEmptyString(payload?.resumeFromRunId);
if (!resumeFromRunId) return null;
const resumeRun = await db
.select({
id: heartbeatRuns.id,
contextSnapshot: heartbeatRuns.contextSnapshot,
sessionIdBefore: heartbeatRuns.sessionIdBefore,
sessionIdAfter: heartbeatRuns.sessionIdAfter,
})
.from(heartbeatRuns)
.where(
and(
eq(heartbeatRuns.id, resumeFromRunId),
eq(heartbeatRuns.companyId, agent.companyId),
eq(heartbeatRuns.agentId, agent.id),
),
)
.then((rows) => rows[0] ?? null);
if (!resumeRun) return null;
const resumeContext = parseObject(resumeRun.contextSnapshot);
const resumeTaskKey = deriveTaskKey(resumeContext, null) ?? taskKey;
const resumeTaskSession = resumeTaskKey
? await getTaskSession(agent.companyId, agent.id, agent.adapterType, resumeTaskKey)
: null;
const sessionCodec = getAdapterSessionCodec(agent.adapterType);
const sessionOverride = buildExplicitResumeSessionOverride({
resumeFromRunId,
resumeRunSessionIdBefore: resumeRun.sessionIdBefore,
resumeRunSessionIdAfter: resumeRun.sessionIdAfter,
taskSession: resumeTaskSession,
sessionCodec,
});
if (!sessionOverride) return null;
return {
resumeFromRunId,
taskKey: resumeTaskKey,
issueId: readNonEmptyString(resumeContext.issueId),
taskId: readNonEmptyString(resumeContext.taskId) ?? readNonEmptyString(resumeContext.issueId),
sessionDisplayId: sessionOverride.sessionDisplayId,
sessionParams: sessionOverride.sessionParams,
};
}
async function resolveWorkspaceForRun(
agent: typeof agents.$inferSelect,
context: Record<string, unknown>,
@ -1920,9 +2017,18 @@ export function heartbeatService(db: Db) {
const resetTaskSession = shouldResetTaskSessionForWake(context);
const sessionResetReason = describeSessionResetReason(context);
const taskSessionForRun = resetTaskSession ? null : taskSession;
const previousSessionParams = normalizeSessionParams(
sessionCodec.deserialize(taskSessionForRun?.sessionParamsJson ?? null),
const explicitResumeSessionParams = normalizeSessionParams(
sessionCodec.deserialize(parseObject(context.resumeSessionParams)),
);
const explicitResumeSessionDisplayId = truncateDisplayId(
readNonEmptyString(context.resumeSessionDisplayId) ??
(sessionCodec.getDisplayId ? sessionCodec.getDisplayId(explicitResumeSessionParams) : null) ??
readNonEmptyString(explicitResumeSessionParams?.sessionId),
);
const previousSessionParams =
explicitResumeSessionParams ??
(explicitResumeSessionDisplayId ? { sessionId: explicitResumeSessionDisplayId } : null) ??
normalizeSessionParams(sessionCodec.deserialize(taskSessionForRun?.sessionParamsJson ?? null));
const config = parseObject(agent.adapterConfig);
const executionWorkspaceMode = resolveExecutionWorkspaceMode({
projectPolicy: projectExecutionWorkspacePolicy,
@ -2098,11 +2204,29 @@ export function heartbeatService(db: Db) {
cleanupReason: null,
});
}
if (issueId && persistedExecutionWorkspace && issueRef?.executionWorkspaceId !== persistedExecutionWorkspace.id) {
await issuesSvc.update(issueId, {
executionWorkspaceId: persistedExecutionWorkspace.id,
...(resolvedProjectWorkspaceId ? { projectWorkspaceId: resolvedProjectWorkspaceId } : {}),
});
if (issueId && persistedExecutionWorkspace) {
const nextIssueWorkspaceMode = issueExecutionWorkspaceModeForPersistedWorkspace(persistedExecutionWorkspace.mode);
const shouldSwitchIssueToExistingWorkspace =
issueRef?.executionWorkspacePreference === "reuse_existing" ||
executionWorkspaceMode === "isolated_workspace" ||
executionWorkspaceMode === "operator_branch";
const nextIssuePatch: Record<string, unknown> = {};
if (issueRef?.executionWorkspaceId !== persistedExecutionWorkspace.id) {
nextIssuePatch.executionWorkspaceId = persistedExecutionWorkspace.id;
}
if (resolvedProjectWorkspaceId && issueRef?.projectWorkspaceId !== resolvedProjectWorkspaceId) {
nextIssuePatch.projectWorkspaceId = resolvedProjectWorkspaceId;
}
if (shouldSwitchIssueToExistingWorkspace) {
nextIssuePatch.executionWorkspacePreference = "reuse_existing";
nextIssuePatch.executionWorkspaceSettings = {
...(issueExecutionWorkspaceSettings ?? {}),
mode: nextIssueWorkspaceMode,
};
}
if (Object.keys(nextIssuePatch).length > 0) {
await issuesSvc.update(issueId, nextIssuePatch);
}
}
if (persistedExecutionWorkspace) {
context.executionWorkspaceId = persistedExecutionWorkspace.id;
@ -2171,7 +2295,8 @@ export function heartbeatService(db: Db) {
}
const runtimeSessionFallback = taskKey || resetTaskSession ? null : runtime.sessionId;
let previousSessionDisplayId = truncateDisplayId(
taskSessionForRun?.sessionDisplayId ??
explicitResumeSessionDisplayId ??
taskSessionForRun?.sessionDisplayId ??
(sessionCodec.getDisplayId ? sessionCodec.getDisplayId(runtimeSessionParams) : null) ??
readNonEmptyString(runtimeSessionParams?.sessionId) ??
runtimeSessionFallback,
@ -2782,7 +2907,9 @@ export function heartbeatService(db: Db) {
payload: promotedPayload,
});
const sessionBefore = await resolveSessionBeforeForWakeup(deferredAgent, promotedTaskKey);
const sessionBefore =
readNonEmptyString(promotedContextSnapshot.resumeSessionDisplayId) ??
await resolveSessionBeforeForWakeup(deferredAgent, promotedTaskKey);
const now = new Date();
const newRun = await tx
.insert(heartbeatRuns)
@ -2861,10 +2988,30 @@ export function heartbeatService(db: Db) {
triggerDetail,
payload,
});
const issueId = readNonEmptyString(enrichedContextSnapshot.issueId) ?? issueIdFromPayload;
let issueId = readNonEmptyString(enrichedContextSnapshot.issueId) ?? issueIdFromPayload;
const agent = await getAgent(agentId);
if (!agent) throw notFound("Agent not found");
const explicitResumeSession = await resolveExplicitResumeSessionOverride(agent, payload, taskKey);
if (explicitResumeSession) {
enrichedContextSnapshot.resumeFromRunId = explicitResumeSession.resumeFromRunId;
enrichedContextSnapshot.resumeSessionDisplayId = explicitResumeSession.sessionDisplayId;
enrichedContextSnapshot.resumeSessionParams = explicitResumeSession.sessionParams;
if (!readNonEmptyString(enrichedContextSnapshot.issueId) && explicitResumeSession.issueId) {
enrichedContextSnapshot.issueId = explicitResumeSession.issueId;
}
if (!readNonEmptyString(enrichedContextSnapshot.taskId) && explicitResumeSession.taskId) {
enrichedContextSnapshot.taskId = explicitResumeSession.taskId;
}
if (!readNonEmptyString(enrichedContextSnapshot.taskKey) && explicitResumeSession.taskKey) {
enrichedContextSnapshot.taskKey = explicitResumeSession.taskKey;
}
issueId = readNonEmptyString(enrichedContextSnapshot.issueId) ?? issueId;
}
const effectiveTaskKey = readNonEmptyString(enrichedContextSnapshot.taskKey) ?? taskKey;
const sessionBefore =
explicitResumeSession?.sessionDisplayId ??
await resolveSessionBeforeForWakeup(agent, effectiveTaskKey);
const writeSkippedRequest = async (skipReason: string) => {
await db.insert(agentWakeupRequests).values({
@ -2928,7 +3075,6 @@ export function heartbeatService(db: Db) {
if (issueId && !bypassIssueExecutionLock) {
const agentNameKey = normalizeAgentNameKey(agent.name);
const sessionBefore = await resolveSessionBeforeForWakeup(agent, taskKey);
const outcome = await db.transaction(async (tx) => {
await tx.execute(
@ -3279,8 +3425,6 @@ export function heartbeatService(db: Db) {
.returning()
.then((rows) => rows[0]);
const sessionBefore = await resolveSessionBeforeForWakeup(agent, taskKey);
const newRun = await db
.insert(heartbeatRuns)
.values({

View file

@ -1,6 +1,7 @@
import { and, asc, desc, eq, inArray, isNull, ne, or, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
activityLog,
agents,
assets,
companies,
@ -62,6 +63,7 @@ function applyStatusSideEffects(
export interface IssueFilters {
status?: string;
assigneeAgentId?: string;
participantAgentId?: string;
assigneeUserId?: string;
touchedByUserId?: string;
unreadForUserId?: string;
@ -134,6 +136,30 @@ function touchedByUserCondition(companyId: string, userId: string) {
`;
}
function participatedByAgentCondition(companyId: string, agentId: string) {
return sql<boolean>`
(
${issues.createdByAgentId} = ${agentId}
OR ${issues.assigneeAgentId} = ${agentId}
OR EXISTS (
SELECT 1
FROM ${issueComments}
WHERE ${issueComments.issueId} = ${issues.id}
AND ${issueComments.companyId} = ${companyId}
AND ${issueComments.authorAgentId} = ${agentId}
)
OR EXISTS (
SELECT 1
FROM ${activityLog}
WHERE ${activityLog.companyId} = ${companyId}
AND ${activityLog.entityType} = 'issue'
AND ${activityLog.entityId} = ${issues.id}::text
AND ${activityLog.agentId} = ${agentId}
)
)
`;
}
function myLastCommentAtExpr(companyId: string, userId: string) {
return sql<Date | null>`
(
@ -508,6 +534,9 @@ export function issueService(db: Db) {
if (filters?.assigneeAgentId) {
conditions.push(eq(issues.assigneeAgentId, filters.assigneeAgentId));
}
if (filters?.participantAgentId) {
conditions.push(participatedByAgentCondition(companyId, filters.participantAgentId));
}
if (filters?.assigneeUserId) {
conditions.push(eq(issues.assigneeUserId, filters.assigneeUserId));
}