@@ -45,8 +45,11 @@ def close_pty_tail(
     collection. Whatever is still waiting then has no later window to complete it, so it is
     replaced here rather than leaving with the session.
 
-    The result is truncated again because the tail is added after the window already applied
-    ``max_output_tokens``, and the returned count covers the text including it.
+    A window that already hit ``max_output_tokens`` is left alone. Its output is at the cap, so
+    the tail sits past it like the rest of what was dropped, and ``original_token_count`` has
+    already told the caller the output is short. Folding it in there would truncate the text a
+    second time and recount the shortened display instead of the source, reporting fewer tokens
+    than the window measured.
     """
     if not leftover:
         return output, original_token_count
@@ -55,9 +58,13 @@ def close_pty_tail(
     if not tail:
         return output, original_token_count
 
+    if original_token_count is not None:
+        return output, original_token_count
+
+    # Nothing was truncated, so this really is the whole output and the count still fits it.
     text = output.decode("utf-8", errors="replace") + tail
-    truncated, recounted = truncate_text_by_tokens(text, max_output_tokens)
-    return truncated.encode("utf-8", errors="replace"), recounted
+    truncated, counted = truncate_text_by_tokens(text, max_output_tokens)
+    return truncated.encode("utf-8", errors="replace"), counted
 
 
 async def flush_pty_tail(
@@ -95,29 +102,38 @@ async def collect_pty_output(
     deadline = time.monotonic() + (yield_time_ms / 1000)
     output = bytearray()
 
-    while True:
-        async with output_lock:
-            while output_chunks:
-                output.extend(output_chunks.popleft())
-
-        if time.monotonic() >= deadline:
-            break
-
-        if is_done():
+    try:
+        while True:
             async with output_lock:
                 while output_chunks:
                     output.extend(output_chunks.popleft())
-            break
-
-        remaining_s = deadline - time.monotonic()
-        if remaining_s <= 0:
-            break
 
-        try:
-            await asyncio.wait_for(output_notify.wait(), timeout=remaining_s)
-        except asyncio.TimeoutError:
-            break
-        output_notify.clear()
+            if time.monotonic() >= deadline:
+                break
+
+            if is_done():
+                async with output_lock:
+                    while output_chunks:
+                        output.extend(output_chunks.popleft())
+                break
+
+            remaining_s = deadline - time.monotonic()
+            if remaining_s <= 0:
+                break
+
+            try:
+                await asyncio.wait_for(output_notify.wait(), timeout=remaining_s)
+            except asyncio.TimeoutError:
+                break
+            output_notify.clear()
+    except asyncio.CancelledError:
+        # Everything drained so far lives only in this buffer, and the session outlives a
+        # cancelled call, so put it back before the cancellation goes on. Otherwise the next
+        # window reads a continuation whose lead byte went with the abandoned call and reports
+        # a replacement character for output that did arrive.
+        if output:
+            output_chunks.appendleft(bytes(output))
+        raise
 
     # Output is collected in repeated windows over one persistent deque, so a character
     # whose bytes straddle a window boundary used to be replaced twice and lost. An
 
@@ -11,6 +11,7 @@
     collect_pty_output,
     flush_pty_tail,
 )
+from agents.sandbox.session.pty_types import truncate_text_by_tokens
 
 
 @pytest.mark.asyncio
@@ -292,3 +293,73 @@ async def window(yield_time_ms: int) -> bytes:
     chunks.append(raw[1:])
     done["value"] = True
     assert (await window(60)).decode("utf-8") == "\u00e9"
+
+
+@pytest.mark.asyncio
+async def test_collect_pty_output_puts_a_drained_window_back_when_cancelled() -> None:
+    # the window drains the lead byte a previous one requeued, then the call is cancelled while
+    # it waits. the session lives on, so those bytes have to go back or its next read reports a
+    # replacement character for output that did arrive
+    raw = "\u00e9".encode()
+    chunks: deque[bytes] = deque([raw[:1]])
+    lock = asyncio.Lock()
+    notify = asyncio.Event()
+    done = {"value": False}
+
+    async def window(yield_time_ms: int) -> bytes:
+        notify.set()
+        collected, _ = await collect_pty_output(
+            output_chunks=chunks,
+            output_lock=lock,
+            output_notify=notify,
+            is_done=lambda: done["value"],
+            yield_time_ms=yield_time_ms,
+            max_output_tokens=None,
+        )
+        return collected
+
+    task = asyncio.create_task(window(60_000))
+    await asyncio.sleep(0.05)
+    assert not chunks
+
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+    assert list(chunks) == [raw[:1]]
+
+    chunks.append(raw[1:])
+    done["value"] = True
+    assert (await window(60)).decode("utf-8") == "\u00e9"
+
+
+def test_close_pty_tail_leaves_a_window_that_already_hit_the_cap_alone() -> None:
+    # the window truncated, so its output is at the cap and the count describes the source.
+    # folding a tail in here would truncate a second time and recount the shortened display
+    display, count = truncate_text_by_tokens("a" * 100, 10)
+    assert count is not None
+
+    output, kept = close_pty_tail(
+        leftover="\u00e9".encode()[:1],
+        output=display.encode(),
+        original_token_count=count,
+        max_output_tokens=10,
+    )
+
+    assert output == display.encode()
+    assert kept == count
+
+
+def test_close_pty_tail_still_folds_the_tail_into_an_untruncated_window() -> None:
+    display, count = truncate_text_by_tokens("hi ", 10)
+    assert count is None
+
+    output, recounted = close_pty_tail(
+        leftover="\u00e9".encode()[:1],
+        output=display.encode(),
+        original_token_count=count,
+        max_output_tokens=10,
+    )
+
+    assert output.decode("utf-8") == "hi \ufffd"
+    assert recounted is None