PaperClipAI/server/src/services/chat-completion-delivery.ts

import { notifyDeliveryWork, DELIVERY_QUEUES } from "./delivery-work-notifications.js";
import { notifyNativeResponsePresentation } from "./native-response-presentation-notifications.js";
import { nativeResponsePresentationCandidates } from "./native-runtime/native-response-presentation-query.js";
import { and, asc, desc, eq, inArray, isNull, lte, sql } from "drizzle-orm";
import { agents, agentWakeupRequests, chatCompletionDeliveries as deliveries, chatTaskHandoffs as handoffs,
  heartbeatRuns, issueComments, issueDocuments, issueThreadInteractions, issues, type Db } from "@paperclipai/db";
import { instanceSettingsService } from "./instance-settings.js";
import { conflict, forbidden, notFound } from "../errors.js";

export const CHAT_COMPLETION_WAKE_REASON = "chat_task_completed";
const MAX_ATTEMPTS = 5;
const LEASE_MS = 60_000;
type Issue = typeof issues.$inferSelect;
type Run = typeof heartbeatRuns.$inferSelect;
type Tx = Parameters<Parameters<Db["transaction"]>[0]>[0];
type Connection = Db | Tx;
const pending = ["pending", "queued"] as const;
const activeRuns = ["queued", "scheduled_retry", "running"];
function ids(run: Run): string[] {
  const value = run.contextSnapshot?.chatCompletionDeliveryIds;
  return Array.isArray(value) ? value.filter((v): v is string => typeof v === "string") : [];
}

/** Human-confirmed suggestions use the same durable audience as native create_task.
 * The source must be read/authorized by the caller; ownership and generation are
 * rechecked under the creation transaction so reset chats cannot receive new work. */
export async function lockAcceptedChatHandoffSource(tx: Connection, input: {
  companyId: string; conversationId: string; sessionGeneration: number;
  userId?: string | null; agentId?: string | null;
}) {
  const [source] = await tx.select().from(issues).where(and(eq(issues.id, input.conversationId),
    eq(issues.companyId, input.companyId))).for("update");
  if (!source?.conversationAgentId) throw notFound("Originating conversation not found");
  if (source.conversationSessionGeneration !== input.sessionGeneration) {
    throw conflict("The originating conversation was reset; request a new task suggestion");
  }
  if (input.userId ? source.conversationUserId !== input.userId : source.conversationAgentId !== input.agentId) {
    throw forbidden("Task handoff requires the originating conversation's participant");
  }
  return source;
}

export async function recordAcceptedChatHandoff(tx: Connection, task: Issue, source: Issue) {
  if (!source.conversationAgentId || source.companyId !== task.companyId || task.conversationAgentId) {
    throw forbidden("Task handoff requires an execution task in the originating conversation's company");
  }
  await tx.insert(handoffs).values({ taskId: task.id, companyId: task.companyId,
    conversationId: source.id, agentId: source.conversationAgentId,
    sessionGeneration: source.conversationSessionGeneration }).onConflictDoNothing();
}

export async function recordChatHandoff(tx: Connection, task: Issue, actorRunId: string | null | undefined) {
  if (!actorRunId || task.conversationAgentId || !task.createdByAgentId) return;
  const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, actorRunId),
    eq(heartbeatRuns.companyId, task.companyId), eq(heartbeatRuns.agentId, task.createdByAgentId)));
  const sourceId = run?.nativeIssueId ?? run?.contextSnapshot?.issueId;
  if (typeof sourceId !== "string" || sourceId === task.id) return;
  const [source] = await tx.select().from(issues).where(and(eq(issues.id, sourceId), eq(issues.companyId, task.companyId)));
  if (!source?.conversationAgentId || source.conversationAgentId !== run.agentId ||
    source.conversationSessionGeneration !== run.contextSnapshot?.conversationSessionGeneration) return;
  await recordAcceptedChatHandoff(tx, task, source);
}

/** Must run on the same transaction as status projection (including native arbitration). */
export async function recordChatCompletion(tx: Connection, before: Issue, after: Issue): Promise<boolean> {
  if (before.status === after.status) return false;
  await tx.update(deliveries).set({ status: "superseded" }).where(and(eq(deliveries.taskId, after.id),
    eq(deliveries.companyId, after.companyId), inArray(deliveries.status, [...pending])));
  if (after.status !== "done") return false;
  const [handoff] = await tx.select().from(handoffs).where(and(eq(handoffs.taskId, after.id), eq(handoffs.companyId, after.companyId)));
  if (!handoff) return false;
  await notifyDeliveryWork(tx, DELIVERY_QUEUES.chatCompletion);
  await tx.insert(deliveries).values({ companyId: after.companyId, taskId: after.id, statusVersion: after.statusVersion }).onConflictDoNothing();
  return true;
}

async function loadAudience(tx: Connection, deliveryId: string) {
  const [row] = await tx.select({ delivery: deliveries, handoff: handoffs, task: issues }).from(deliveries)
    .innerJoin(handoffs, and(eq(handoffs.taskId, deliveries.taskId), eq(handoffs.companyId, deliveries.companyId)))
    .innerJoin(issues, and(eq(issues.id, deliveries.taskId), eq(issues.companyId, deliveries.companyId)))
    .where(eq(deliveries.id, deliveryId));
  if (!row) return null;
  const [source] = await tx.select().from(issues).where(and(eq(issues.id, row.handoff.conversationId), eq(issues.companyId, row.handoff.companyId)));
  return { ...row, source };
}
function current(row: NonNullable<Awaited<ReturnType<typeof loadAudience>>>) {
  return row.task.status === "done" && !["superseded", "exhausted"].includes(row.delivery.status) &&
    row.source?.conversationAgentId === row.handoff.agentId &&
    row.source.conversationSessionGeneration === row.handoff.sessionGeneration;
}

async function taskResult(tx: Connection, task: Issue) {
  // Never copy worker-authored titles, comments or document bodies into the
  // source agent's instructions. These links and lifecycle facts are generated
  // by the server; the saved work remains on its normal access-controlled task.
  const documents = await tx.select({ id: issueDocuments.documentId }).from(issueDocuments)
    .where(and(eq(issueDocuments.companyId, task.companyId), eq(issueDocuments.issueId, task.id))).limit(1);
  // A bounded receipt projection gives the coordinator actual verdicts to link,
  // rather than asking it to infer whether the worker's prose needs another review.
  // These are historical decisions on their named versions, not publication authority.
  const reviews = await tx.select().from(issueThreadInteractions).where(and(
    eq(issueThreadInteractions.companyId, task.companyId), eq(issueThreadInteractions.issueId, task.id),
    inArray(issueThreadInteractions.kind, ["request_confirmation", "request_item_verdicts"]),
    inArray(issueThreadInteractions.status, ["accepted", "answered"]),
  )).orderBy(desc(issueThreadInteractions.resolvedAt), desc(issueThreadInteractions.id)).limit(20);
  const recordedOutputDecisions = reviews.flatMap(review => {
    const payload = review.payload as Record<string, any>, result = (review.result ?? {}) as Record<string, any>;
    const candidates = review.kind === "request_confirmation" ? [{ target: payload.target, outcome: "approve", human: Boolean(review.resolvedByUserId) && !review.resolvedByAgentId }]
      : (Array.isArray(result.items) ? result.items : []).map((item: any) => ({
        target: (Array.isArray(payload.items) ? payload.items : []).find((entry: any) => entry.id === item.id)?.target ?? payload.target,
        outcome: item.verdict, human: Boolean(item.resolvedByUserId) && !item.resolvedByAgentId,
      }));
    return candidates.filter(candidate => candidate.human && ["approve", "reject", "defer"].includes(candidate.outcome) &&
      candidate.target && ["issue_document", "issue_attachment"].includes(candidate.target.type) &&
      (!candidate.target.issueId || candidate.target.issueId === task.id) && candidate.target.key !== "plan")
      .map(candidate => ({ interactionId: review.id, outcome: candidate.outcome,
        target: candidate.target.type === "issue_document"
          ? { type: "issue_document", key: candidate.target.key, revisionId: candidate.target.revisionId }
          : { type: "issue_attachment", attachmentId: candidate.target.attachmentId, assetId: candidate.target.assetId, sha256: candidate.target.sha256 },
        url: `/issues/${task.identifier ?? task.id}#interaction-${review.id}` }));
  }).slice(0, 30);
  return { id: task.id, identifier: task.identifier, status: task.status,
    completedAt: task.completedAt, url: `/issues/${task.identifier ?? task.id}`,
    hasSavedDocuments: documents.length > 0, recordedOutputDecisions };
}

/** A completed onboarding parent still owes the result of its own child handoff.
 * This permits a reporting turn without reopening Done or reviving cancellation. */
export async function isCompletedOnboardingHandoffWake(db: Connection, input: {
  companyId: string; issueId: string; agentId: string; reason: string | null;
  contextSnapshot: Record<string, unknown>;
}) {
  if (input.reason !== "issue_children_completed" || typeof input.contextSnapshot.completedChildIssueId !== "string") return false;
  const [source] = await db.select().from(issues).where(and(eq(issues.id, input.issueId), eq(issues.companyId, input.companyId)));
  if (source?.originKind !== "onboarding_first_task" || source.status !== "done" || source.assigneeAgentId !== input.agentId) return false;
  const children = await db.select({ id: issues.id, status: issues.status }).from(issues)
    .where(and(eq(issues.companyId, input.companyId), eq(issues.parentId, source.id)));
  return children.some(child => child.id === input.contextSnapshot.completedChildIssueId && child.status === "done") &&
    children.every(child => ["done", "cancelled"].includes(child.status));
}

/** Freeze the input at turn start. New completions cannot be consumed by an already-running turn. */
export async function prepareChatCompletionTurn(db: Db, run: Run): Promise<Run> {
  const issueId = run.contextSnapshot?.issueId;
  if (typeof issueId !== "string") return run;
  return db.transaction(async tx => {
    const [source] = await tx.select().from(issues).where(and(eq(issues.id, issueId), eq(issues.companyId, run.companyId))).for("update");
    if (!source) return run;
    if (source.originKind === "onboarding_first_task" && run.contextSnapshot?.wakeReason === "issue_children_completed") {
      const children = await tx.select().from(issues).where(and(eq(issues.companyId, run.companyId), eq(issues.parentId, source.id), eq(issues.status, "done"))).orderBy(issues.createdAt).limit(20);
      const contextSnapshot = { ...run.contextSnapshot, onboardingCompletion: true,
        chatCompletionUpdates: await Promise.all(children.map(child => taskResult(tx, child))) };
      await tx.update(heartbeatRuns).set({ contextSnapshot }).where(eq(heartbeatRuns.id, run.id));
      return { ...run, contextSnapshot };
    }
    const initialIds = ids(run);
    // Only completion wakes absorb pending events; ordinary user turns retain their existing input.
    if (initialIds.length === 0 && run.contextSnapshot?.wakeReason !== CHAT_COMPLETION_WAKE_REASON) return run;
    if (run.contextSnapshot?.conversationSessionGeneration !== source.conversationSessionGeneration) throw new Error("chat_completion_superseded");
    const rows = await tx.select({ delivery: deliveries }).from(deliveries).innerJoin(handoffs, eq(handoffs.taskId, deliveries.taskId))
      .where(and(eq(deliveries.companyId, run.companyId), eq(handoffs.conversationId, issueId),
        eq(handoffs.agentId, run.agentId), eq(handoffs.sessionGeneration, source.conversationSessionGeneration),
        inArray(deliveries.status, [...pending]),
        sql`(${deliveries.targetRunId} is null or ${deliveries.targetRunId} = ${run.id})`))
      .orderBy(asc(deliveries.createdAt)).limit(20);
    const updates: Awaited<ReturnType<typeof taskResult>>[] = [];
    const accepted: string[] = [];
    for (const { delivery } of rows) {
      const row = await loadAudience(tx, delivery.id);
      if (!row || !current(row)) {
        await tx.update(deliveries).set({ status: "superseded" }).where(eq(deliveries.id, delivery.id));
        continue;
      }
      accepted.push(delivery.id); updates.push(await taskResult(tx, row.task));
      await tx.update(deliveries).set({ status: "queued", targetRunId: run.id }).where(eq(deliveries.id, delivery.id));
    }
    if (accepted.length === 0) throw new Error("chat_completion_superseded");
    const contextSnapshot = { ...run.contextSnapshot, chatCompletionDeliveryIds: accepted, chatCompletionUpdates: updates };
    await tx.update(heartbeatRuns).set({ contextSnapshot }).where(eq(heartbeatRuns.id, run.id));
    return { ...run, contextSnapshot };
  });
}

export function chatCompletionInstruction(context: Record<string, unknown>) {
  if (!Array.isArray(context.chatCompletionUpdates) || !context.chatCompletionUpdates.length) return "";
  return `\n\nDelegated work has completed. Tell the user in this conversation what finished and provide access using the supplied task links. Report completion only for the tasks listed in this update; other tasks receive their own completion updates. Use the recorded status and result locations; do not repeat a promise to do work that is already Done. Do not start more work, change these tasks, or create duplicate review requests. Reuse and link the task’s existing canonical review/decision; completion updates are notifications, not fresh quality-review authority. recordedOutputDecisions are human verdicts on the exact versions shown. Link them instead of re-requesting review of those versions; do not treat approval as permission to publish. A separate explicit user request may authorize a new review. Keep opaque IDs and hashes out of your visible reply. The following JSON contains server-recorded lifecycle facts and result locations:\n${JSON.stringify(context.chatCompletionUpdates)}`;
}

/** Called under the comment transaction, before insertion. An event can publish only once. */
export async function existingChatCompletionReply(tx: Connection, runId: string, issueId: string) {
  const [run] = await tx.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
  if (!run || ids(run).length === 0) return null;
  const rows = await tx.select().from(deliveries).where(and(eq(deliveries.companyId, run.companyId), inArray(deliveries.id, ids(run))));
  for (const delivery of rows) {
    // Match the status writer's task-first lock order, then reload the outbox
    // row so a concurrent reopen cannot leave us with its pre-commit snapshot.
    await tx.select({ id: issues.id }).from(issues).where(and(eq(issues.id, delivery.taskId), eq(issues.companyId, run.companyId))).for("update");
    const row = await loadAudience(tx, delivery.id);
    if (!row || row.source?.id !== issueId || row.handoff.agentId !== run.agentId || !current(row)) throw new Error("chat_completion_superseded");
    if (row.delivery.responseCommentId) {
      const [comment] = await tx.select().from(issueComments).where(and(eq(issueComments.id, row.delivery.responseCommentId), eq(issueComments.issueId, issueId)));
      if (comment) return comment;
    }
    if (row.delivery.status !== "queued" || row.delivery.targetRunId !== runId) throw new Error("chat_completion_superseded");
  }
  if (rows.length !== ids(run).length) throw new Error("chat_completion_superseded");
  return null;
}
export async function acknowledgeChatCompletionReply(tx: Connection, runId: string, commentId: string) {
  await tx.update(deliveries).set({ status: "delivered", responseCommentId: commentId, error: null })
    .where(and(eq(deliveries.targetRunId, runId), eq(deliveries.status, "queued")));
}

export function chatCompletionDeliveryService(db: Db, heartbeat: { wakeup(agentId: string, options: {
  source: "automation"; triggerDetail: "system"; reason: string; idempotencyKey: string; allowRunCoalescing: boolean;
  requestedByActorType: "system"; requestedByActorId: string; payload: Record<string, unknown>; contextSnapshot: Record<string, unknown>;
}): Promise<{ id: string } | null> }) {
  async function deliver(id: string): Promise<void> {
    // Lease outbox work before leaving the transaction; recovery reuses the same wake key.
    const claimed = await db.update(deliveries).set({ nextAttemptAt: new Date(Date.now() + LEASE_MS) })
      .where(and(eq(deliveries.id, id), inArray(deliveries.status, [...pending]), lte(deliveries.nextAttemptAt, new Date())))
      .returning().then(rows => rows[0]);
    if (!claimed) return;
    try {
      const row = await loadAudience(db, id);
      if (!row || !current(row)) {
        await db.update(deliveries).set({ status: "superseded" }).where(eq(deliveries.id, id)); return;
      }
      if (!(await instanceSettingsService(db).getExperimental()).enableAgentChat) return;
      const [agent] = await db.select().from(agents).where(and(eq(agents.id, row.handoff.agentId), eq(agents.companyId, claimed.companyId)));
      if (!agent || ["paused", "terminated"].includes(agent.status)) return;
      const key = `chat-completion:${id}:${claimed.attempts}`;
      const wakes = await db.select().from(agentWakeupRequests).where(and(eq(agentWakeupRequests.companyId, claimed.companyId), eq(agentWakeupRequests.idempotencyKey, key)));
      const runId = claimed.targetRunId ?? wakes.find(w => w.runId)?.runId;
      if (runId) {
        const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
        if (run && activeRuns.includes(run.status)) return;

        const presentationDeadline = run?.runtimeMode === "native" && run.status === "succeeded" && run.finishedAt
          ? run.finishedAt.getTime() + MAX_ATTEMPTS * LEASE_MS : null;
        if (run && presentationDeadline !== null && Date.now() < presentationDeadline) {
          const scope = { companyId: claimed.companyId, issueId: row.handoff.conversationId,
            runId: run.id, agentId: row.handoff.agentId };
          const deferred = await db.transaction(async tx => {


            const candidates = nativeResponsePresentationCandidates(scope).as("pending_native_completion_reply");
            const [response] = await tx.select().from(candidates);
            if (!response || Date.now() >= presentationDeadline) return false;
            await notifyNativeResponsePresentation(tx, scope);
            const updated = await tx.update(deliveries).set({ nextAttemptAt: new Date(Math.min(presentationDeadline, Math.max(
              claimed.nextAttemptAt.getTime(), response.nextAttemptAt?.getTime() ?? 0,
            ))) }).where(and(eq(deliveries.id, id), eq(deliveries.companyId, claimed.companyId),
              eq(deliveries.status, "queued"), eq(deliveries.targetRunId, run.id),
              eq(deliveries.statusVersion, claimed.statusVersion), eq(deliveries.attempts, claimed.attempts),
              eq(deliveries.nextAttemptAt, claimed.nextAttemptAt))).returning({ id: deliveries.id });
            return updated.length > 0;
          });
          if (deferred) return;
        }
        const advanced = await db.transaction(async tx => {


          const [task] = await tx.select({ status: issues.status }).from(issues)
            .where(and(eq(issues.id, claimed.taskId), eq(issues.companyId, claimed.companyId))).for("update");
          if (task?.status !== "done") return false;
          const sameClaim = and(eq(deliveries.id, id), eq(deliveries.companyId, claimed.companyId),
            inArray(deliveries.status, [...pending]), eq(deliveries.statusVersion, claimed.statusVersion),
            eq(deliveries.attempts, claimed.attempts),
            eq(deliveries.nextAttemptAt, claimed.nextAttemptAt),
            claimed.targetRunId ? eq(deliveries.targetRunId, claimed.targetRunId) : isNull(deliveries.targetRunId));
          const [latest] = await tx.select({ id: deliveries.id }).from(deliveries).where(sameClaim).for("update");
          if (!latest) return false;
          await notifyDeliveryWork(tx, DELIVERY_QUEUES.chatCompletion);
          const updated = await tx.update(deliveries).set({ status: claimed.attempts + 1 >= MAX_ATTEMPTS ? "exhausted" : "pending",
            attempts: claimed.attempts + 1, targetRunId: null, nextAttemptAt: new Date(),
            error: run?.error ?? "Completion turn ended without a reply" }).where(sameClaim).returning({ id: deliveries.id });
          return updated.length > 0;
        });


        if (advanced && claimed.attempts + 1 < MAX_ATTEMPTS) await deliver(id);
        return;
      }
      if (wakes.some(w => ["queued", "claimed", "deferred_issue_execution", "coalesced"].includes(w.status))) return;
      if (wakes.length) {
        // Terminal receipts must advance the key; reusing them collides forever.
        const dailyCap = wakes.some(w => w.reason?.startsWith("heartbeat.daily_"));
        const next = new Date(Date.now() + LEASE_MS);
        if (dailyCap) next.setUTCHours(24, 0, 0, 0);
        await db.update(deliveries).set({ attempts: claimed.attempts + 1, targetRunId: null,
          status: !dailyCap && claimed.attempts + 1 >= MAX_ATTEMPTS ? "exhausted" : "pending",
          nextAttemptAt: next, error: wakes[0].reason ?? "Completion wake did not start" })
          .where(and(eq(deliveries.id, id), inArray(deliveries.status, [...pending])));
        return;
      }
      // One undispatched head per conversation. Its turn absorbs the remaining
      // events at admission; events arriving after admission need a later turn.
      const siblings = await db.select({ delivery: deliveries, run: heartbeatRuns }).from(deliveries)
        .innerJoin(handoffs, eq(handoffs.taskId, deliveries.taskId))
        .leftJoin(heartbeatRuns, eq(heartbeatRuns.id, deliveries.targetRunId))
        .where(and(eq(deliveries.companyId, claimed.companyId), eq(handoffs.conversationId, row.handoff.conversationId),
          eq(handoffs.sessionGeneration, row.handoff.sessionGeneration), inArray(deliveries.status, [...pending])))
        .orderBy(asc(deliveries.createdAt), asc(deliveries.id));
      if (siblings.some(s => s.run && activeRuns.includes(s.run.status) &&
        !Array.isArray(s.run.contextSnapshot?.chatCompletionUpdates))) return;
      if (siblings.find(s => !s.delivery.targetRunId)?.delivery.id !== id) return;
      const run = await heartbeat.wakeup(row.handoff.agentId, {
        source: "automation", triggerDetail: "system", reason: CHAT_COMPLETION_WAKE_REASON, idempotencyKey: key, allowRunCoalescing: false,
        requestedByActorType: "system", requestedByActorId: "chat_completion_delivery",
        payload: { issueId: row.handoff.conversationId, chatCompletionDeliveryIds: [id] },
        contextSnapshot: { issueId: row.handoff.conversationId, taskId: row.handoff.conversationId,
          wakeReason: CHAT_COMPLETION_WAKE_REASON, chatCompletionDeliveryIds: [id], conversationSessionGeneration: row.handoff.sessionGeneration },
      });
      if (!run) {
        const receipts = await db.select({ id: agentWakeupRequests.id }).from(agentWakeupRequests)
          .where(and(eq(agentWakeupRequests.companyId, claimed.companyId), eq(agentWakeupRequests.idempotencyKey, key)));
        if (!receipts.length) throw new Error("Completion wake returned no run or durable receipt");
      }
      if (run) await db.update(deliveries).set({ targetRunId: run.id, status: "queued" })
        .where(and(eq(deliveries.id, id), inArray(deliveries.status, [...pending]), sql`${deliveries.targetRunId} is null`));
    } catch (error) {
      const receipts = await db.select({ id: agentWakeupRequests.id }).from(agentWakeupRequests)
        .where(and(eq(agentWakeupRequests.companyId, claimed.companyId), eq(agentWakeupRequests.idempotencyKey, `chat-completion:${id}:${claimed.attempts}`)));
      await db.update(deliveries).set({ error: error instanceof Error ? error.message : String(error),
        ...(receipts.length ? {} : { attempts: claimed.attempts + 1,
          status: claimed.attempts + 1 >= MAX_ATTEMPTS ? "exhausted" as const : "pending" as const }) })
        .where(and(eq(deliveries.id, id), inArray(deliveries.status, [...pending])));
    }
  }
  async function sweepPending(scope?: { companyId: string; taskId: string }) {
    const due = await db.select({ id: deliveries.id }).from(deliveries)
      .where(and(inArray(deliveries.status, [...pending]), lte(deliveries.nextAttemptAt, new Date()),
        ...(scope ? [eq(deliveries.companyId, scope.companyId), eq(deliveries.taskId, scope.taskId)] : [])))
      .orderBy(asc(deliveries.nextAttemptAt)).limit(100);
    for (const row of due) await deliver(row.id);
  }
  async function hasPending() {
    return (await db.select({ id: deliveries.id }).from(deliveries)
      .where(inArray(deliveries.status, [...pending])).limit(1)).length > 0;
  }
  return { deliver, sweepPending, hasPending };
}