From cc8df64e46a1c26410598911dd89e20eb60fe12b Mon Sep 17 00:00:00 2001 From: JulioMCruz Date: Sat, 18 Jul 2026 15:55:00 -0400 Subject: [PATCH] keep workflow output safe and canonical --- src/tools/createDoc.ts | 32 +++++++++++++++++++++++++++++++- src/tools/outputSanitizer.ts | 12 ++++++++++++ src/tools/postDocMessage.ts | 3 ++- src/tools/postProjectMessage.ts | 19 +++++++++++++++---- src/tools/proposePlan.ts | 8 +++++++- src/tools/updateTaskStatus.ts | 28 +++++++++++++++++++++++++++- tests/tools-pure.test.ts | 11 +++++++++++ 7 files changed, 105 insertions(+), 8 deletions(-) create mode 100644 src/tools/outputSanitizer.ts diff --git a/src/tools/createDoc.ts b/src/tools/createDoc.ts index dfe124b..8b7d56b 100644 --- a/src/tools/createDoc.ts +++ b/src/tools/createDoc.ts @@ -51,8 +51,38 @@ export const createDoc: Tool = { }; } + // Approved workflows have one canonical document. Workers may ask to + // create a task-specific spec, but returning the approved plan keeps every + // contribution in the shared workspace and makes retries idempotent. + const projectData = project.data() as Record; + const workflow = (projectData.workflow ?? null) as + | { phase?: string; planId?: string } + | null; + if ( + workflow?.planId && + ["running", "pm_review", "complete"].includes(String(workflow.phase)) + ) { + const canonicalRef = projectRef.collection("docs").doc(workflow.planId); + const canonical = await canonicalRef.get(); + if (canonical.exists) { + return { + ok: true, + data: { + docId: workflow.planId, + type: canonical.data()?.type ?? "plan", + draft: canonical.data()?.draft === true, + activePlan: true, + canonical: true, + reused: true, + }, + message: + "This approved workflow uses one canonical project document. Reuse this docId instead of creating a separate deliverable doc.", + }; + } + } + if (args.type === "plan") { - const activePlanId = project.data()?.activePlanId; + const activePlanId = projectData.activePlanId; if (typeof activePlanId === "string" && activePlanId.length > 0) { const activePlanRef = projectRef.collection("docs").doc(activePlanId); const activePlan = await activePlanRef.get(); diff --git a/src/tools/outputSanitizer.ts b/src/tools/outputSanitizer.ts new file mode 100644 index 0000000..96dc6c8 --- /dev/null +++ b/src/tools/outputSanitizer.ts @@ -0,0 +1,12 @@ +/** Credentials used to claim work must never become part of user content. */ +export function redactClaimTokens(value: string): string { + return value + .replace( + /\bclaimToken\s*[:=]\s*["']?[A-Za-z0-9_-]{16,128}["']?/gi, + "claimToken=[redacted]", + ) + .replace( + /\bclaim token\s*[:=]\s*["']?[A-Za-z0-9_-]{16,128}["']?/gi, + "claim token=[redacted]", + ); +} diff --git a/src/tools/postDocMessage.ts b/src/tools/postDocMessage.ts index fa10da6..82a040d 100644 --- a/src/tools/postDocMessage.ts +++ b/src/tools/postDocMessage.ts @@ -15,6 +15,7 @@ import { FieldValue } from "firebase-admin/firestore"; import { db } from "../firestore.js"; import { docIdSchema, projectIdSchema } from "./docShared.js"; import type { Tool } from "./types.js"; +import { redactClaimTokens } from "./outputSanitizer.js"; const InputSchema = z .object({ @@ -49,7 +50,7 @@ export const postDocMessage: Tool = { const msgRef = docRef.collection("messages").doc(); await msgRef.set({ from: "agent", - text: args.text, + text: redactClaimTokens(args.text), agentName: ctx.convId ?? "agent", createdAt: FieldValue.serverTimestamp(), }); diff --git a/src/tools/postProjectMessage.ts b/src/tools/postProjectMessage.ts index 9048c0b..5702795 100644 --- a/src/tools/postProjectMessage.ts +++ b/src/tools/postProjectMessage.ts @@ -13,6 +13,7 @@ import { z } from "zod"; import { db } from "../firestore.js"; import { postProjectChat } from "../projectChat.js"; import type { Tool } from "./types.js"; +import { redactClaimTokens } from "./outputSanitizer.js"; const InputSchema = z .object({ @@ -46,18 +47,28 @@ export const postProjectMessage: Tool = { message: `No project "${args.projectId}" for this wallet.`, }; } + const projectData = project.data() as Record; + const workflow = (projectData.workflow ?? null) as + | { phase?: string; convId?: string } + | null; + const workflowConvId = + workflow?.convId && ["running", "pm_review"].includes(String(workflow.phase)) + ? workflow.convId + : undefined; const message = await postProjectChat({ wallet: ctx.wallet, projectId: args.projectId, - convId: (project.data()?.chatConvId as string | undefined) ?? undefined, + convId: + workflowConvId ?? + ((projectData.chatConvId as string | undefined) ?? undefined), sender: ctx.convId, - text: args.text, + text: redactClaimTokens(args.text), targets: Array.from( new Set( [ `user:${ctx.wallet}`, - project.data()?.pmAgent - ? `agent:${String(project.data()?.pmAgent)}` + projectData.pmAgent + ? `agent:${String(projectData.pmAgent)}` : null, ].filter((identity): identity is string => Boolean(identity)), ), diff --git a/src/tools/proposePlan.ts b/src/tools/proposePlan.ts index 64b8032..b384523 100644 --- a/src/tools/proposePlan.ts +++ b/src/tools/proposePlan.ts @@ -72,6 +72,11 @@ export const proposePlan: Tool = { const now = FieldValue.serverTimestamp(); const projectRef = docRef.parent.parent!; + const project = await projectRef.get(); + const workflowConvId = + typeof project.data()?.chatConvId === "string" + ? (project.data()?.chatConvId as string) + : undefined; const batch = projectRef.firestore.batch(); batch.set( docRef, @@ -85,6 +90,7 @@ export const proposePlan: Tool = { phase: "awaiting_approval", planId: refs.docId, taskIds: [], + ...(workflowConvId ? { convId: workflowConvId } : {}), updatedAt: now, }, updatedAt: now, @@ -111,7 +117,7 @@ export const proposePlan: Tool = { const chatDelivered = await postProjectChat({ wallet: ctx.wallet, projectId: args.projectId, - convId: ((await projectRef.get()).data()?.chatConvId as string | undefined) ?? undefined, + convId: workflowConvId, sender: ctx.convId, targets: [`user:${ctx.wallet}`], text: `Plan proposed with ${tasks} task${tasks === 1 ? "" : "s"}. Review the plan and approve it before work starts.`, diff --git a/src/tools/updateTaskStatus.ts b/src/tools/updateTaskStatus.ts index 5c8b86f..77d9e00 100644 --- a/src/tools/updateTaskStatus.ts +++ b/src/tools/updateTaskStatus.ts @@ -32,6 +32,7 @@ import { FieldValue } from "firebase-admin/firestore"; import { db } from "../firestore.js"; import { logActivity } from "../activityEvents.js"; import type { Tool } from "./types.js"; +import { redactClaimTokens } from "./outputSanitizer.js"; const ProofSchema = z .object({ @@ -134,7 +135,7 @@ export const updateTaskStatus: Tool = { updatedAt: FieldValue.serverTimestamp(), lastWorkerUpdateAt: FieldValue.serverTimestamp(), }; - if (typeof args.result === "string") patch.result = args.result; + if (typeof args.result === "string") patch.result = redactClaimTokens(args.result); if (args.proof && args.proof.length > 0) { patch.proof = FieldValue.arrayUnion( ...args.proof.map((p) => ({ ...p, at: Date.now() })), @@ -163,6 +164,31 @@ export const updateTaskStatus: Tool = { await taskRef.update(patch); + // Keep the dispatcher-wide per-agent lease in sync with the task claim. + // This is what preserves one active task per single-session runtime across + // API replicas and across projects, including tasks longer than 10 min. + if (claimActive && args.claimToken === claim?.token && data.agent?.trim()) { + const leaseRef = db() + .collection("agent_dispatch_leases") + .doc(`${ctx.wallet}__${data.agent.trim()}`); + await db().runTransaction(async (tx) => { + const lease = await tx.get(leaseRef); + if (!lease.exists || lease.data()?.token !== args.claimToken) return; + if (effectiveStatus === "Done" || effectiveStatus === "Review") { + tx.delete(leaseRef); + return; + } + tx.set( + leaseRef, + { + expiresAtMs: Date.now() + 10 * 60_000, + updatedAt: FieldValue.serverTimestamp(), + }, + { merge: true }, + ); + }); + } + // Activity feed: narrate the transition in plain language (only on a real // status CHANGE — claim-heartbeat updates with the same status are noise). if (effectiveStatus !== current) { diff --git a/tests/tools-pure.test.ts b/tests/tools-pure.test.ts index c0298e5..a7de25b 100644 --- a/tests/tools-pure.test.ts +++ b/tests/tools-pure.test.ts @@ -29,6 +29,7 @@ process.env.AUDIT_ENABLED = "false"; const { getRunbookFor } = await import("../src/tools/getRunbookFor.js"); const { searchKnowledge } = await import("../src/tools/searchKnowledge.js"); const { explainPlugin } = await import("../src/tools/explainPlugin.js"); +const { redactClaimTokens } = await import("../src/tools/outputSanitizer.js"); const ctx = { wallet: "0x" + "a".repeat(40), @@ -133,3 +134,13 @@ describe("explainPlugin", () => { if (!r.ok) expect(r.errorClass).toBe("NOT_FOUND"); }); }); + +describe("redactClaimTokens", () => { + it("removes claim credentials from persisted user content", () => { + expect( + redactClaimTokens( + "done — roadmap delivered; claimToken=617b2ada-bf27-4665-8393-c31d47d1e151", + ), + ).toBe("done — roadmap delivered; claimToken=[redacted]"); + }); +});