Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions src/__tests__/unit/lib/proposals/codeChangeCompletion.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ vi.mock("@/services/swarm/createPr", async (importOriginal) => {
});

import {
attachPrArtifact,
completeClaimFromResult,
markClaimRunFailed,
reconcileClaim,
Expand Down Expand Up @@ -272,3 +273,45 @@ describe("reconcileClaim", () => {
expect(mockReconcilePr).not.toHaveBeenCalled();
});
});

describe("attachPrArtifact", () => {
const PR_URL = "https://github.com/acme/widgets/pull/7";
const REPO_URL = "https://github.com/acme/widgets";

it("targets the ASSISTANT message so the PR lands beside its diff", async () => {
// A claim Task seeds two messages in one transaction — the USER prompt and
// the ASSISTANT diff — so their createdAt can tie. Ordering alone could
// hand back the prompt row; the role filter is what makes this decidable.
vi.mocked(db.chatMessage.findFirst).mockResolvedValue({
id: "msg-assistant",
artifacts: [],
} as never);

await attachPrArtifact(TASK_ID, PR_URL, REPO_URL);

expect(db.chatMessage.findFirst).toHaveBeenCalledWith(
expect.objectContaining({
where: { taskId: TASK_ID, role: "ASSISTANT" },
}),
);
expect(db.artifact.create).toHaveBeenCalledWith(
expect.objectContaining({
data: expect.objectContaining({
messageId: "msg-assistant",
type: "PULL_REQUEST",
}),
}),
);
});

it("does not duplicate an artifact the webhook already attached", async () => {
vi.mocked(db.chatMessage.findFirst).mockResolvedValue({
id: "msg-assistant",
artifacts: [{ id: "existing" }],
} as never);

await attachPrArtifact(TASK_ID, PR_URL, REPO_URL);

expect(db.artifact.create).not.toHaveBeenCalled();
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -541,3 +541,84 @@ describe("approveCodeChange — reconcile on retry", () => {
expect(mockReconcilePr).not.toHaveBeenCalled();
});
});

describe("approveCodeChange — claim Task seeding", () => {
/**
* Swap in a transaction mock whose `tx` client is retained, so the rows the
* claim transaction writes can be inspected after the call.
*/
function captureTx() {
const tx = {
task: {
create: vi.fn().mockResolvedValue({ id: CLAIM_TASK_ID }),
update: vi.fn().mockResolvedValue({}),
},
chatMessage: { create: vi.fn().mockResolvedValue({ id: SEED_MSG_ID }) },
artifact: {
create: vi.fn().mockResolvedValue({}),
updateMany: vi.fn().mockResolvedValue({}),
},
};
vi.mocked(db.$transaction).mockImplementation(async (arg: unknown) => {
if (typeof arg !== "function") return undefined as never;
return (arg as (t: unknown) => Promise<unknown>)(tx) as never;
});
return tx;
}

function outputWithPrompt(prompt: string) {
const base = codeChangeOutput();
return { ...base, payload: { ...base.payload, prompt } };
}

it("creates the claim IN_PROGRESS so the tasks list does not hide it", async () => {
const tx = captureTx();
mockFetchStored.mockResolvedValue([msg(codeChangeOutput())]);

await approve([msg(codeChangeOutput())]);

// A TODO claim matches none of the list's visibility branches — it is
// neither agent-mode nor Stakwork-backed — so it would surface only in
// Kanban. COMPLETED must survive alongside it for pr-monitor's fix path.
expect(tx.task.create).toHaveBeenCalledWith(
expect.objectContaining({
data: expect.objectContaining({
status: "IN_PROGRESS",
workflowStatus: "COMPLETED",
}),
}),
);
});

it("seeds the originating prompt as a USER message ahead of the diff", async () => {
const tx = captureTx();
mockFetchStored.mockResolvedValue([
msg(outputWithPrompt("Add a null check in auth middleware")),
]);

// The prompt is read off the STORED proposal, never the caller's copy.
await approve([msg(outputWithPrompt("rm -rf /"))]);

expect(tx.chatMessage.create).toHaveBeenCalledTimes(2);

const [first, second] = tx.chatMessage.create.mock.calls;
expect(first[0].data).toMatchObject({
role: "USER",
message: "Add a null check in auth middleware",
});
// The diff rides on the ASSISTANT message, which `attachPrArtifact` later
// resolves by role to hang the PULL_REQUEST artifact off.
expect(second[0].data.role).toBe("ASSISTANT");
expect(second[0].data.artifacts.create.type).toBe("DIFF");
});

it("seeds only the diff message for a proposal stored before `prompt` existed", async () => {
const tx = captureTx();
mockFetchStored.mockResolvedValue([msg(codeChangeOutput())]);

await approve([msg(codeChangeOutput())]);

expect(tx.chatMessage.create).toHaveBeenCalledTimes(1);
expect(tx.chatMessage.create.mock.calls[0][0].data.role).toBe("ASSISTANT");
});
});
6 changes: 3 additions & 3 deletions src/__tests__/unit/services/release-stale-task-pods.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ describe("releaseStaleTaskPods", () => {
deleted: false,
OR: [
{ podId: { not: null } },
{ status: "IN_PROGRESS", workflowStatus: { not: "HALTED" } },
{ status: "IN_PROGRESS", workflowStatus: { not: "HALTED" }, proposalId: null },
],
},
select: {
Expand Down Expand Up @@ -307,10 +307,10 @@ describe("releaseStaleTaskPods", () => {
const findManyCall = vi.mocked(mockDb.task.findMany).mock.calls[2][0];
// Should use OR clause to find both:
// 1. Tasks with pods (any status)
// 2. Stale IN_PROGRESS tasks without pods
// 2. Stale IN_PROGRESS tasks without pods, excluding code-change claims
expect(findManyCall?.where?.OR).toEqual([
{ podId: { not: null } },
{ status: "IN_PROGRESS", workflowStatus: { not: "HALTED" } },
{ status: "IN_PROGRESS", workflowStatus: { not: "HALTED" }, proposalId: null },
]);

vi.useRealTimers();
Expand Down
1 change: 1 addition & 0 deletions src/lib/ai/codeChangeTools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -438,6 +438,7 @@ export function buildCodeChangeTools(ctx: CapabilityContext): ToolSet {
diffSha256,
filesChanged: diffs.length,
baseBranchDisplay,
prompt,
},
// Stamped server-side from the tool's own context — never a caller
// input. The approval handler reads this back off the STORED
Expand Down
2 changes: 1 addition & 1 deletion src/lib/constants/prompt.ts
Original file line number Diff line number Diff line change
Expand Up @@ -408,7 +408,7 @@ You **cannot** create Workspaces or Repositories — for those, tell the user to

**You are not a coding agent.** You never make code changes yourself. \`repo_agent\` is **STRICTLY READ-ONLY investigation** — it has no write path and cannot open a PR. Two proposal tools can produce real code changes, routed by scope:

- **\`propose_code_change\`** — use for **single-repo, single-concern, well-scoped changes** when this tool is in your tool list. It works in workspaces with any number of repositories: \`repositoryUrl\` names the one you are patching, and it must be registered in that workspace. What has to be true is that the CHANGE lands in one repo — not that the workspace only owns one. Hard requirements (enforced server-side): the change must not touch \`prisma/schema.prisma\` or any migration file, and the diff must be ≤ 50 files / 200 KB. The tool generates a real diff preview — the user reviews it and clicks Approve to open the PR under their own GitHub identity. Do NOT use for: changes spanning multiple repos, schema/migration files, large refactors, or when you cannot tell which repo the change belongs in.
- **\`propose_code_change\`** — use for **single-repo, single-concern, well-scoped changes** when this tool is in your tool list. It works in workspaces with any number of repositories: \`repositoryUrl\` names the one you are patching, and it must be registered in that workspace. What has to be true is that the CHANGE lands in one repo — not that the workspace only owns one. Keep the change small and focused — aim for < 10 files. Hard requirements (enforced server-side): the change must not touch \`prisma/schema.prisma\` or any migration file, and the diff must be ≤ 50 files / 200 KB. The tool generates a real diff preview — the user reviews it and clicks Approve to open the PR under their own GitHub identity. Do NOT use for: changes spanning multiple repos, schema/migration files, large refactors, or when you cannot tell which repo the change belongs in.
- **\`propose_feature\`** — use for **everything else**: work spanning multiple repos, any schema/migration, over the file/byte cap, ambiguous scope, or when \`propose_code_change\` is not available. The feature pipeline routes the work to the appropriate coding agent.

The boundary for the hard caps (schema path, size) is mechanical. Repo choice and "ambiguous scope" are judgment-based: pick \`repositoryUrl\` from where the code you are changing actually lives — you normally know it already, because you found the file with \`repo_agent\` or the user named it. If you would be guessing between repos, that itself is the signal to use \`propose_feature\`. Never pick a repo just because it is the workspace's first one.
Expand Down
9 changes: 7 additions & 2 deletions src/lib/proposals/codeChangeCompletion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
*/

import { db } from "@/lib/db";
import { ArtifactType } from "@prisma/client";
import { ArtifactType, ChatRole } from "@prisma/client";
import {
_processCompletedResult,
reconcilePr,
Expand Down Expand Up @@ -104,6 +104,11 @@ export function parseCreatePrClaim(value: unknown): CreatePrClaim | null {
* Best-effort and idempotent: a failure here must never turn a landed PR
* into a reported failure, and a concurrent webhook/reconcile must not
* produce two artifacts for the same PR.
*
* Scoped to the ASSISTANT message — the one carrying the DIFF. A claim Task
* also seeds a USER message holding the originating prompt, and both rows are
* written in the same transaction, so `createdAt` alone can tie and resolve
* either way. The role filter keeps the PR landing beside its diff.
*/
export async function attachPrArtifact(
taskId: string,
Expand All @@ -112,7 +117,7 @@ export async function attachPrArtifact(
): Promise<void> {
try {
const msg = await db.chatMessage.findFirst({
where: { taskId },
where: { taskId, role: ChatRole.ASSISTANT },
orderBy: { createdAt: "desc" },
select: {
id: true,
Expand Down
22 changes: 22 additions & 0 deletions src/lib/proposals/handleApproval.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ import {
ArtifactType,
WorkflowStatus,
TaskSourceType,
TaskStatus,
} from "@prisma/client";
import type { PullRequestContent, DiffContent } from "@/lib/chat";
import { mcpCreatePrompt, mcpUpdatePrompt } from "@/lib/mcp/mcpTools";
Expand Down Expand Up @@ -751,6 +752,12 @@ async function approveCodeChange(args: {
updatedById: userId,
sourceType: TaskSourceType.SYSTEM,
mode: "live",
// The tasks list hides TODO tasks that are neither agent-mode nor
// Stakwork-backed, so at the TODO default the claim only ever showed
// in Kanban. IN_PROGRESS is also the honest state: the PR run is in
// flight here. pr-monitor / the GitHub webhook carry it to DONE on
// merge (or CANCELLED on close) off the PULL_REQUEST artifact below.
status: TaskStatus.IN_PROGRESS,
workflowStatus: WorkflowStatus.COMPLETED,
stakworkProjectId: null,
podId: null,
Expand All @@ -775,6 +782,21 @@ async function approveCodeChange(args: {
const diffDiffs = unifiedDiffToActionResults(payload.diff, repoName);
const diffContent: DiffContent = { diffs: diffDiffs };

// Seed the originating instruction first, so the Task view reads as a
// request followed by its result. Absent on proposals stored before
// `prompt` joined the payload — those render result-only.
if (payload.prompt) {
await tx.chatMessage.create({
data: {
taskId: task.id,
message: payload.prompt,
role: ChatRole.USER,
status: ChatStatus.SENT,
},
select: { id: true },
});
}

const msg = await tx.chatMessage.create({
data: {
taskId: task.id,
Expand Down
6 changes: 6 additions & 0 deletions src/lib/proposals/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,12 @@ export interface CodeChangeProposalPayload {
filesChanged: number;
/** Display-only: the base branch the swarm resolved server-side. */
baseBranchDisplay?: string;
/** The instruction sent to the read-only `repo_agent` run that produced
* `diff`. Display-only: seeded as a USER-role ChatMessage on the claim Task
* so the Task view shows what was asked, not just the result. Never re-sent
* to the swarm — approval forwards `diff`. Absent on proposals stored
* before this field existed; those replay result-only. */
prompt?: string;
}

/** What the propose tools return from `execute(...)` on success. */
Expand Down
9 changes: 8 additions & 1 deletion src/services/task-coordinator-cron.ts
Original file line number Diff line number Diff line change
Expand Up @@ -437,10 +437,17 @@ export async function releaseStaleTaskPods(): Promise<{
OR: [
// Tasks with pods (any status) - release pod
{ podId: { not: null } },
// IN_PROGRESS tasks without pods - just halt
// IN_PROGRESS tasks without pods - just halt.
// Code-change claims are exempt: they are created IN_PROGRESS with
// workflowStatus COMPLETED (which pr-monitor's fix path needs) and
// own no pod, so halting would only clobber that status in the one
// case worth preserving - a PR that never reported back, which the
// code-change-reconcile cron still expects to recover. proposalId is
// written solely by approveCodeChange, so it identifies them exactly.
{
status: "IN_PROGRESS",
workflowStatus: { not: "HALTED" },
proposalId: null,
},
],
},
Expand Down
Loading