Stop refunding consumed Agent recovery usage

app/components/FinishReasonNotice.tsx

@@ -21,11 +21,8 @@ export const FinishReasonNotice = ({
   finishReason,
   onContinue,
 }: FinishReasonNoticeProps) => {
-  const { dataStream, isAutoResuming, isAutoContinuing } = useDataStreamState();
+  const { isAutoResuming, isAutoContinuing } = useDataStreamState();
   const [hasContinued, setHasContinued] = useState(false);
-  const usageProtected = dataStream.some(
-    (part) => part.type === "data-auto-continue-usage-protected",
-  );
 
   if (!finishReason) return null;
 
@@ -44,28 +41,39 @@ export const FinishReasonNotice = ({
 
   const getNoticeContent = () => {
     if (finishReason === "tool-calls") {
-      return <>Reached the step limit for this turn.</>;
+      return (
+        <>Reached the step limit for this turn. Completed work was saved.</>
+      );
     }
 
     if (finishReason === "timeout" || finishReason === "preemptive-timeout") {
-      return <>Reached the time limit for this turn.</>;
+      return (
+        <>Reached the time limit for this turn. Completed work was saved.</>
+      );
     }
 
     if (finishReason === OUTPUT_LIMIT_FINISH_REASON) {
       return (
         <>
           The response reached its output limit before finishing. Continue to
-          resume where it stopped.
+          resume where it stopped; completed work was saved.
         </>
       );
     }
 
     if (finishReason === "context-limit") {
-      return <>Reached the context limit for this conversation.</>;
+      return (
+        <>
+          Reached the context limit for this conversation. Completed work was
+          saved.
+        </>
+      );
     }
 
     if (finishReason === POST_SUMMARIZATION_INCOMPLETE_FINISH_REASON) {
-      return <>Paused after compacting the conversation.</>;
+      return (
+        <>Paused after compacting the conversation. Completed work was saved.</>
+      );
     }
 
     if (finishReason === AGENT_RUN_SPEND_CAP_FINISH_REASON) {
@@ -93,15 +101,7 @@ export const FinishReasonNotice = ({
   return (
     <div className="mt-2 w-full">
       <div className="bg-muted text-muted-foreground rounded-lg px-3 py-2 border border-border flex items-center justify-between gap-3 flex-wrap">
-        <div className="flex flex-col gap-1">
-          <span>{content}</span>
-          {usageProtected ? (
-            <span className="text-foreground">
-              This automatic recovery attempt didn&apos;t finish, so its usage
-              was restored.
-            </span>
-          ) : null}
-        </div>
+        <span>{content}</span>
         {showContinue && (
           <Button
             type="button"

app/components/__tests__/FinishReasonNotice.test.tsx

@@ -8,40 +8,30 @@ import { DataStreamProvider, useDataStream } from "../DataStreamProvider";
 import { MAX_AUTO_CONTINUES } from "@/app/hooks/useAutoContinue";
 import { POST_SUMMARIZATION_INCOMPLETE_FINISH_REASON } from "@/lib/chat/stop-conditions";
 import type { ChatMode, SelectedModel } from "@/types/chat";
-import type { ScopedDataUIPart } from "../DataStreamProvider";
 
 function DataStreamSetter({
   isAutoResuming,
   isAutoContinuing,
   autoContinueCount,
-  dataStream,
   children,
 }: {
   isAutoResuming?: boolean;
   isAutoContinuing?: boolean;
   autoContinueCount?: number;
-  dataStream?: ScopedDataUIPart[];
   children: React.ReactNode;
 }) {
-  const {
-    setDataStream,
-    setIsAutoResuming,
-    setIsAutoContinuing,
-    setAutoContinueCount,
-  } = useDataStream();
+  const { setIsAutoResuming, setIsAutoContinuing, setAutoContinueCount } =
+    useDataStream();
 
   React.useEffect(() => {
     if (isAutoResuming !== undefined) setIsAutoResuming(isAutoResuming);
     if (isAutoContinuing !== undefined) setIsAutoContinuing(isAutoContinuing);
     if (autoContinueCount !== undefined)
       setAutoContinueCount(autoContinueCount);
-    if (dataStream !== undefined) setDataStream(dataStream);
   }, [
     isAutoResuming,
     isAutoContinuing,
     autoContinueCount,
-    dataStream,
-    setDataStream,
     setIsAutoResuming,
     setIsAutoContinuing,
     setAutoContinueCount,
@@ -63,7 +53,6 @@ function renderNotice(
     isAutoResuming?: boolean;
     isAutoContinuing?: boolean;
     autoContinueCount?: number;
-    dataStream?: ScopedDataUIPart[];
   },
 ) {
   return render(
@@ -175,26 +164,10 @@ describe("FinishReasonNotice", () => {
       ).toBeInTheDocument();
     });
 
-    it("confirms when an incomplete automatic recovery had its usage restored", () => {
-      renderNotice(
-        { finishReason: "tool-calls", mode: "agent" },
-        {
-          isAutoResuming: false,
-          autoContinueCount: MAX_AUTO_CONTINUES,
-          dataStream: [
-            {
-              type: "data-auto-continue-usage-protected",
-              data: { status: "restored" },
-            },
-          ],
-        },
-      );
+    it("confirms that completed work was preserved at the step limit", () => {
+      renderNotice({ finishReason: "tool-calls", mode: "agent" });
 
-      expect(
-        screen.getByText(
-          /This automatic recovery attempt didn't finish, so its usage was restored/i,
-        ),
-      ).toBeInTheDocument();
+      expect(screen.getByText(/Completed work was saved/i)).toBeInTheDocument();
     });
 
     it.each([

app/components/worked-for-parts.ts

@@ -9,7 +9,6 @@ const TRAILING_METADATA_PART_TYPES = new Set([
   "data-agent-auto-review-lifecycle",
   "data-appendMessage",
   "data-auto-continue",
-  "data-auto-continue-usage-protected",
   "data-context-usage",
   "data-diff",
   "data-file-metadata",

lib/api/__tests__/agent-long-contracts.test.ts

@@ -1198,7 +1198,7 @@ describe("agent-long task — Trigger.dev dashboard error visibility", () => {
       /const genericDelegationFlagPromise\s*=\s*agentPermissionMode === "full_access"/,
     );
     expect(routeSrc).toMatch(
-      /existingChat,[\s\S]{0,200}userCustomization,[\s\S]{0,200}genericDelegationEnabled,[\s\S]{0,200}Promise\.all/,
+      /existingChat, userCustomization, genericDelegationEnabled[\s\S]*Promise\.all/,
     );
     expect(routeSrc).toContain('"agent-generic-delegation-v1"');
     expect(routeSrc).not.toContain("securityValidationSubagentsEnabled");
@@ -1415,62 +1415,51 @@ describe("agent-long task — Trigger.dev dashboard error visibility", () => {
       expect(source).toMatch(/getAgentAutoContinueStopSource\(\{/);
       expect(source).toMatch(/finishReason:\s*state\.streamFinishReason/);
       expect(source).toMatch(
-        /if \(\s*autoContinueStopSource\s*&&[\s\S]{0,100}!isAutomaticContinuation\s*\)/,
+        /if \(\s*autoContinueStopSource[\s\S]{0,100}!isAutomaticContinuation\s*\) \{\s*writeAutoContinue/,
       );
       expect(source).toMatch(/writeAutoContinue\(writer\)/);
       expect(source).toMatch(/agent_auto_continue_suppressed/);
     }
   });
 
-  test("both Agent backends protect only a bounded automatic recovery that stops incomplete", () => {
+  test("both Agent backends require an explicit continuation after elapsed timeout", () => {
     for (const source of [chatHandlerSrc, taskSrc]) {
-      expect(source).toMatch(
-        /if \(isAutomaticContinuation\) \{\s*await protectIncompleteAutomaticContinuation\(\{/,
+      const autoContinueIdx = source.lastIndexOf(
+        "const autoContinueStopSource =",
       );
-      expect(source).toMatch(
-        /assignment:\s*autoContinueUsageProtectionAssignment/,
+      const autoContinueEndIdx = source.indexOf("});", autoContinueIdx);
+
+      expect(autoContinueIdx).toBeGreaterThan(-1);
+      expect(autoContinueEndIdx).toBeGreaterThan(autoContinueIdx);
+      expect(source.slice(autoContinueIdx, autoContinueEndIdx)).not.toContain(
+        "stoppedDueToElapsedTimeout",
       );
-      expect(source).toMatch(/stopSource,/);
     }
-    expect(routeSrc).toMatch(/AGENT_AUTO_CONTINUE_USAGE_PROTECTION_FLAG/);
-    expect(routeSrc).toMatch(/autoContinueUsageProtectionAssignment,/);
-    expect(taskSrc).toMatch(
-      /autoContinueUsageProtectionAssignment\?:\s*AgentAutoContinueUsageProtectionAssignment/,
-    );
   });
 
-  test("fallback and final retry completion paths preserve automatic recovery protection", () => {
-    const directRetryDeductionIdx = chatHandlerSrc.indexOf(
-      "await deductAccumulatedUsage(retryMessageId)",
-    );
-    const directRetryProtectionIdx = chatHandlerSrc.indexOf(
-      "await protectTerminalAutomaticContinuation()",
-      directRetryDeductionIdx,
-    );
-    expect(directRetryDeductionIdx).toBeGreaterThan(-1);
-    expect(directRetryProtectionIdx).toBeGreaterThan(directRetryDeductionIdx);
-
-    const triggerFinalizerIdx = taskSrc.indexOf(
-      "const finalizeRetryStream = async",
-    );
-    const triggerRetryDeductionIdx = taskSrc.indexOf(
-      "await deductAccumulatedUsage()",
-      triggerFinalizerIdx,
-    );
-    const triggerRetryProtectionIdx = taskSrc.indexOf(
-      "await protectTerminalAutomaticContinuation()",
-      triggerRetryDeductionIdx,
-    );
-    expect(triggerFinalizerIdx).toBeGreaterThan(-1);
-    expect(triggerRetryProtectionIdx).toBeGreaterThan(triggerRetryDeductionIdx);
+  test("incomplete Agent turns persist and bill observed work before offering continuation", () => {
+    for (const source of [chatHandlerSrc, taskSrc]) {
+      const autoContinueIdx = source.lastIndexOf(
+        "const autoContinueStopSource =",
+      );
+      const persistenceIdx = source.lastIndexOf(
+        "sendFileMetadataToStream(accumulatedFiles)",
+        autoContinueIdx,
+      );
+      const deductionIdx = source.lastIndexOf(
+        "await deductAccumulatedUsage",
+        autoContinueIdx,
+      );
 
-    const finalRetryFinishIdx = taskSrc.indexOf("messages: finalRetryMessages");
-    const finalRetryFinalizeIdx = taskSrc.indexOf(
-      "await finalizeRetryStream({",
-      finalRetryFinishIdx,
-    );
-    expect(finalRetryFinishIdx).toBeGreaterThan(-1);
-    expect(finalRetryFinalizeIdx).toBeGreaterThan(finalRetryFinishIdx);
+      expect(autoContinueIdx).toBeGreaterThan(-1);
+      expect(persistenceIdx).toBeGreaterThan(-1);
+      expect(deductionIdx).toBeGreaterThan(-1);
+      expect(persistenceIdx).toBeLessThan(autoContinueIdx);
+      expect(deductionIdx).toBeLessThan(autoContinueIdx);
+      expect(source.slice(deductionIdx, autoContinueIdx)).not.toContain(
+        "usageRefundTracker.refund",
+      );
+    }
   });
 
   test("the direct chat boundary strictly normalizes automatic continuation flags", () => {
@@ -1680,7 +1669,7 @@ describe("agent-long task — Trigger.dev dashboard error visibility", () => {
       expect(source).toMatch(/getAgentAutoContinueStopSource\(\{/);
       expect(source).toMatch(/autoContinueStopSource/);
       expect(source).toMatch(
-        /if \(\s*autoContinueStopSource\s*&&[\s\S]{0,100}!isAutomaticContinuation\s*\)/,
+        /if \(\s*autoContinueStopSource[\s\S]{0,100}!isAutomaticContinuation\s*\) \{\s*writeAutoContinue/,
       );
       expect(source).toMatch(/agent_auto_continue_signaled/);
     }