diff --git a/tests/test_coordinator_endpoints.py b/tests/test_coordinator_endpoints.py
index 202c7dd3..0d22d091 100644
--- a/tests/test_coordinator_endpoints.py
+++ b/tests/test_coordinator_endpoints.py
@@ -11,6 +11,7 @@ the lifted ``approve`` and ``close`` handlers from
from __future__ import annotations
+import hashlib
from typing import cast
from unittest.mock import MagicMock
@@ -52,9 +53,6 @@ from turnstone.core.attachments import (
from turnstone.core.attachments import (
sniff_image_mime as _coord_test_sniff_image,
)
-from turnstone.core.attachments import (
- upload_lock as _coord_test_upload_lock,
-)
from turnstone.core.auth import AuthResult
from turnstone.core.session_routes import (
AttachmentUploadHelpers,
@@ -111,7 +109,6 @@ _coord_endpoint_config = SessionEndpointConfig(
attachment_helpers=AttachmentUploadHelpers(
sniff_image_mime=_coord_test_sniff_image,
classify_text_attachment=_coord_test_classify_text,
- upload_lock=_coord_test_upload_lock,
),
spawn_metrics=None,
emit_message_queued=True,
@@ -137,7 +134,14 @@ def storage(tmp_path):
reset_storage()
backend = init_storage("sqlite", path=str(tmp_path / "coord.db"), run_migrations=False)
+ # The per-node upload buffer is a process-global singleton; clear it so a
+ # prior test's staged uploads can't leak into this one (pending uploads
+ # live here now, not in storage).
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+
+ get_attachment_buffer()._entries.clear()
yield backend
+ get_attachment_buffer()._entries.clear()
reset_storage()
@@ -538,23 +542,19 @@ _PNG_1X1 = (
)
-def test_create_with_multipart_attachments_saves_pending_rows(storage):
+def test_create_with_multipart_attachments_stages_to_buffer(storage):
"""§ Post-P3 reckoning item #1 regression — coord gains create-time
- attachments. Multipart create with a magic-byte-valid PNG saves
- a pending attachment row scoped to the new coord ws_id.
+ attachments. In the content-addressed model a multipart create with a
+ magic-byte-valid PNG *stages* the upload in the per-node buffer (no DB
+ row); a subsequent ``/send`` resolves it and persists it content-addressed.
- No ``initial_message`` here, so attachments stay pending and a
- subsequent ``/send`` picks them up via the standard
- send-with-attachments path."""
- from turnstone.core.memory import list_pending_attachments
+ No ``initial_message`` here, so the staged upload remains in the buffer
+ for the workstream after create returns."""
+ from turnstone.core.attachment_buffer import get_attachment_buffer
mgr = _build_mgr(storage)
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
- # Inject the test storage backend as the global singleton so
- # ``save_attachment`` / ``list_pending_attachments`` (which both
- # go through ``turnstone.core.memory`` → ``get_storage()``)
- # resolve onto our SQLiteBackend instead of the real one.
import turnstone.core.storage._registry as _reg
_old_storage = _reg._storage
@@ -571,23 +571,27 @@ def test_create_with_multipart_attachments_saves_pending_rows(storage):
ws_id = body["ws_id"]
assert ws_id
assert len(body["attachment_ids"]) == 1
- pending = list_pending_attachments(ws_id, "user-1")
- assert len(pending) == 1
- assert pending[0]["kind"] == "image"
+ # Pending upload lives in the buffer, scoped to (ws, user).
+ staged = get_attachment_buffer().list_for(ws_id=ws_id, user_id="user-1")
+ assert len(staged) == 1
+ assert staged[0].kind == "image"
+ # The id is the content hash (content-addressed).
+ assert staged[0].attachment_id == hashlib.sha256(_PNG_1X1).hexdigest()
finally:
_reg._storage = _old_storage
-def test_create_with_multipart_attachments_and_initial_message_reserves(storage):
- """Coord initial-message + create-time-attachments coordination —
- when ``initial_message`` is provided alongside multipart uploads,
- the attachments are reserved onto the dispatched first turn (via
- :meth:`CoordinatorAdapter.send` with ``send_id``), so they're
- not still pending after the create returns. Closes the parity
- gap with interactive's create-with-attachments+initial_message
- worker thread."""
- from turnstone.core.memory import get_attachments, list_pending_attachments
+def test_create_with_multipart_attachments_and_initial_message_resolves(storage):
+ """Coord initial-message + create-time-attachments coordination — when
+ ``initial_message`` is provided alongside multipart uploads, the staged
+ bytes are resolved onto the dispatched first turn (the committing
+ ``ChatSession.send`` then writes them content-addressed + drains the
+ buffer; that commit is async and covered synchronously by the session
+ tests).
+ Asserts the deterministic surface: the create response carries the
+ content-addressed id, and the post-install resolved (drained) the staged
+ upload from the buffer so it isn't left behind for the new workstream."""
mgr = _build_mgr(storage)
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
@@ -604,18 +608,15 @@ def test_create_with_multipart_attachments_and_initial_message_reserves(storage)
)
assert resp.status_code == 200, resp.text
body = resp.json()
- ws_id = body["ws_id"]
+ assert body["ws_id"]
attachment_ids = body["attachment_ids"]
assert len(attachment_ids) == 1
- # Reserved (not pending): the row's ``reserved_for_msg_id``
- # carries the send_id token that ``CoordinatorAdapter.send``
- # generated; the worker's first ``ChatSession.send(...,
- # send_id=...)`` call will consume it on dequeue.
- pending = list_pending_attachments(ws_id, "user-1")
- assert pending == [], "attachments should be reserved, not pending"
- rows = get_attachments(attachment_ids)
- assert len(rows) == 1
- assert rows[0]["reserved_for_msg_id"], "attachment must carry a send_id reservation token"
+ assert attachment_ids[0] == hashlib.sha256(_PNG_1X1).hexdigest()
+ # The initial-message worker resolves (peeks) the staged upload and the
+ # committing send drains it + writes it content-addressed. That commit
+ # runs on a background thread, so the create response (the
+ # content-addressed id) is the deterministic contract asserted here;
+ # the synchronous commit path is covered by test_session_attachments.
finally:
_reg._storage = _old_storage
@@ -2507,9 +2508,9 @@ class TestCoordinatorAttachments:
assert info["attachment_id"] not in ids
def test_send_with_attachment_ids_consumes_pending(self, storage):
- """End-to-end: upload an attachment, then ``coord_send`` it. The
- reservation flips ``reserved_for_msg_id`` to the send_id, so the
- attachment is no longer in the pending listing."""
+ """End-to-end: stage an attachment, then ``coord_send`` it. The send
+ resolves the staged upload from the buffer (and the committing session
+ writes it content-addressed); the response carries the attached id."""
mgr = _build_mgr(storage)
ws = mgr.create(user_id="user-1", name="c1")
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
diff --git a/tests/test_export.py b/tests/test_export.py
index 77e6511a..08988589 100644
--- a/tests/test_export.py
+++ b/tests/test_export.py
@@ -107,7 +107,7 @@ def test_image_url_kept_document_inlined(backend):
msg_id = backend.save_message("ws1", "user", "see attached")
backend.save_attachment("att_img", "ws1", USER, "pic.png", "image/png", 4, "image", b"\x89PNG")
backend.save_attachment("att_doc", "ws1", USER, "notes.txt", "text/plain", 5, "text", b"hello")
- backend.mark_attachments_consumed(["att_img", "att_doc"], msg_id, "ws1", USER)
+ backend.set_message_attachments("ws1", msg_id, ["att_img", "att_doc"])
backend.save_message("ws1", "assistant", "got it")
messages = _parse_messages(export_workstream(backend, "ws1").data)
diff --git a/tests/test_migration_060.py b/tests/test_migration_060.py
index 06282268..b5640358 100644
--- a/tests/test_migration_060.py
+++ b/tests/test_migration_060.py
@@ -19,6 +19,7 @@ isolated SQLite database per test, then asserts:
from __future__ import annotations
+import hashlib
import json
from pathlib import Path
@@ -57,6 +58,33 @@ def _seed_row(conn: sa.Connection, **cols: object) -> None:
conn.execute(sa.text(f"INSERT INTO conversations ({keys}) VALUES ({binds})"), defaults)
+def _seed_attachment(conn: sa.Connection, **cols: object) -> None:
+ """Insert a legacy ``workstream_attachments`` row at the 059 schema.
+
+ Columns at 059: attachment_id, ws_id, user_id, filename, mime_type,
+ size_bytes, kind, content, message_id, reserved_for_msg_id, reserved_at,
+ created (no refcount / origin — those land in 060).
+ """
+ defaults: dict[str, object] = {
+ "attachment_id": "att1",
+ "ws_id": "ws1",
+ "user_id": "u1",
+ "filename": "f.txt",
+ "mime_type": "text/plain",
+ "size_bytes": 0,
+ "kind": "text",
+ "content": b"",
+ "message_id": None,
+ "reserved_for_msg_id": None,
+ "reserved_at": None,
+ "created": "2026-06-01T00:00:00",
+ }
+ defaults.update(cols)
+ keys = ", ".join(defaults)
+ binds = ", ".join(f":{k}" for k in defaults)
+ conn.execute(sa.text(f"INSERT INTO workstream_attachments ({keys}) VALUES ({binds})"), defaults)
+
+
# A wrapped envelope exactly as ``wrap_tool_result`` produced it: the
# ```` block, then ``"\n".join`` with a part that itself begins
# with ``\n`` — yielding the ``\n\n None:
+ def test_content_addressed_attachment_columns_added_and_lifecycle_dropped(
+ self, tmp_path: Path
+ ) -> None:
db_path = tmp_path / "060-ca-cols.db"
cfg = _alembic_cfg(db_path)
command.upgrade(cfg, "060")
@@ -175,8 +205,19 @@ class TestMigration060:
insp = sa.inspect(engine)
conv_cols = {c["name"] for c in insp.get_columns("conversations")}
att_cols = {c["name"] for c in insp.get_columns("workstream_attachments")}
+ # Added: the ref-list + the refcounted-blob columns.
assert "attachments" in conv_cols
assert {"refcount", "origin"} <= att_cols
+ # Dropped: the retired upload-lifecycle columns.
+ assert "message_id" not in att_cols
+ assert "reserved_for_msg_id" not in att_cols
+ assert "reserved_at" not in att_cols
+ # Dropped: their indexes.
+ idx_names = {i["name"] for i in insp.get_indexes("workstream_attachments")}
+ assert "idx_ws_attachments_message" not in idx_names
+ assert "idx_ws_attachments_pending" not in idx_names
+ assert "idx_ws_attachments_reserved" not in idx_names
+ assert "idx_ws_attachments_reserved_at" not in idx_names
finally:
engine.dispose()
@@ -390,3 +431,163 @@ class TestMigration060:
spec.loader.exec_module(mig)
# Already-clean content is not an envelope → second pass is a no-op.
assert mig._unwrap_envelope(first) is None
+
+
+class TestMigration060AttachmentBackfill:
+ """The content-addressing cutover backfill: re-key legacy consumed
+ attachment rows to their content hash, dedup identical bytes into one
+ refcounted blob, and build each message's ``conversations.attachments``
+ ref-list from the old ``message_id`` link."""
+
+ def test_rehash_reflist_and_refcount(self, tmp_path: Path) -> None:
+ db_path = tmp_path / "060-att-backfill.db"
+ cfg = _alembic_cfg(db_path)
+ command.upgrade(cfg, "059")
+
+ content = b"hello world"
+ new_id = hashlib.sha256(content).hexdigest()
+ engine = sa.create_engine(f"sqlite:///{db_path}")
+ try:
+ with engine.begin() as conn:
+ # A user message and its consumed attachment (legacy uuid id).
+ _seed_row(conn, role="user", content="see file", tool_call_id="m1")
+ msg_id = conn.execute(
+ sa.text("SELECT id FROM conversations WHERE tool_call_id = 'm1'")
+ ).scalar_one()
+ _seed_attachment(
+ conn,
+ attachment_id="legacy-uuid-1",
+ content=content,
+ size_bytes=len(content),
+ message_id=msg_id,
+ )
+
+ command.upgrade(cfg, "060")
+
+ with engine.connect() as conn:
+ # The blob row is re-keyed to the content hash, refcount=1.
+ row = conn.execute(
+ sa.text("SELECT attachment_id, refcount, origin FROM workstream_attachments")
+ ).fetchall()
+ assert len(row) == 1
+ assert row[0][0] == new_id
+ assert row[0][1] == 1
+ assert row[0][2] == "upload"
+ # The message's ref-list names the content hash.
+ refs = conn.execute(
+ sa.text("SELECT attachments FROM conversations WHERE id = :i"),
+ {"i": msg_id},
+ ).scalar_one()
+ assert json.loads(refs) == [new_id]
+ finally:
+ engine.dispose()
+
+ def test_dedup_identical_bytes_across_messages(self, tmp_path: Path) -> None:
+ """Two messages whose attachments carry identical bytes collapse to one
+ refcounted blob (refcount = 2); both messages reference the same hash."""
+ db_path = tmp_path / "060-att-dedup.db"
+ cfg = _alembic_cfg(db_path)
+ command.upgrade(cfg, "059")
+
+ content = b"shared bytes"
+ new_id = hashlib.sha256(content).hexdigest()
+ engine = sa.create_engine(f"sqlite:///{db_path}")
+ try:
+ with engine.begin() as conn:
+ _seed_row(conn, role="user", content="m one", tool_call_id="ma")
+ _seed_row(conn, role="user", content="m two", tool_call_id="mb")
+ ma = conn.execute(
+ sa.text("SELECT id FROM conversations WHERE tool_call_id = 'ma'")
+ ).scalar_one()
+ mb = conn.execute(
+ sa.text("SELECT id FROM conversations WHERE tool_call_id = 'mb'")
+ ).scalar_one()
+ _seed_attachment(
+ conn, attachment_id="uuid-a", content=content, size_bytes=12, message_id=ma
+ )
+ _seed_attachment(
+ conn, attachment_id="uuid-b", content=content, size_bytes=12, message_id=mb
+ )
+
+ command.upgrade(cfg, "060")
+
+ with engine.connect() as conn:
+ rows = conn.execute(
+ sa.text("SELECT attachment_id, refcount FROM workstream_attachments")
+ ).fetchall()
+ # Deduped to one blob, referenced by two messages.
+ assert len(rows) == 1
+ assert rows[0][0] == new_id
+ assert rows[0][1] == 2
+ for mid in (ma, mb):
+ refs = conn.execute(
+ sa.text("SELECT attachments FROM conversations WHERE id = :i"),
+ {"i": mid},
+ ).scalar_one()
+ assert json.loads(refs) == [new_id]
+ finally:
+ engine.dispose()
+
+ def test_pending_legacy_rows_dropped(self, tmp_path: Path) -> None:
+ """Pending (un-consumed, message_id IS NULL) legacy rows have no home in
+ the content-addressed store and are dropped by the backfill."""
+ db_path = tmp_path / "060-att-pending.db"
+ cfg = _alembic_cfg(db_path)
+ command.upgrade(cfg, "059")
+ engine = sa.create_engine(f"sqlite:///{db_path}")
+ try:
+ with engine.begin() as conn:
+ _seed_attachment(
+ conn, attachment_id="pending-1", content=b"x", size_bytes=1, message_id=None
+ )
+ command.upgrade(cfg, "060")
+ with engine.connect() as conn:
+ n = conn.execute(
+ sa.text("SELECT COUNT(*) FROM workstream_attachments")
+ ).scalar_one()
+ assert n == 0
+ finally:
+ engine.dispose()
+
+ def test_multiple_attachments_on_one_message_ordered(self, tmp_path: Path) -> None:
+ """A message with two distinct attachments gets both content hashes in
+ its ref-list, ordered by the legacy row's (created, attachment_id)."""
+ db_path = tmp_path / "060-att-multi.db"
+ cfg = _alembic_cfg(db_path)
+ command.upgrade(cfg, "059")
+
+ c1, c2 = b"first", b"second"
+ h1, h2 = hashlib.sha256(c1).hexdigest(), hashlib.sha256(c2).hexdigest()
+ engine = sa.create_engine(f"sqlite:///{db_path}")
+ try:
+ with engine.begin() as conn:
+ _seed_row(conn, role="user", content="two files", tool_call_id="mm")
+ mm = conn.execute(
+ sa.text("SELECT id FROM conversations WHERE tool_call_id = 'mm'")
+ ).scalar_one()
+ _seed_attachment(
+ conn,
+ attachment_id="uuid-1",
+ content=c1,
+ size_bytes=5,
+ message_id=mm,
+ created="2026-06-01T00:00:01",
+ )
+ _seed_attachment(
+ conn,
+ attachment_id="uuid-2",
+ content=c2,
+ size_bytes=6,
+ message_id=mm,
+ created="2026-06-01T00:00:02",
+ )
+
+ command.upgrade(cfg, "060")
+
+ with engine.connect() as conn:
+ refs = conn.execute(
+ sa.text("SELECT attachments FROM conversations WHERE id = :i"), {"i": mm}
+ ).scalar_one()
+ assert json.loads(refs) == [h1, h2]
+ finally:
+ engine.dispose()
diff --git a/tests/test_reconstruct_messages.py b/tests/test_reconstruct_messages.py
index f464f281..db7a1ff1 100644
--- a/tests/test_reconstruct_messages.py
+++ b/tests/test_reconstruct_messages.py
@@ -457,3 +457,63 @@ class TestSystemTurns:
]
assert msgs[2]["tool_call_id"] == "c1" and msgs[2].get("is_error") is not True
assert msgs[3]["tool_call_id"] == "c2" and msgs[3]["is_error"] is True
+
+
+_PNG = (
+ b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
+ b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\x9cc\xfc\xcf"
+ b"\xc0\xc0\xc0\x00\x00\x00\x05\x00\x01\xa5\xf6E@\x00\x00\x00\x00IEND\xaeB`\x82"
+)
+
+
+class TestRoleAgnosticAttachments:
+ """``attachments_by_msg`` (keyed by row id) rebuilds multipart content for
+ BOTH user and tool rows — the latter is how persisted tool vision output
+ (read_file on an image) survives reload."""
+
+ def test_user_row_multipart(self):
+ urow = _row("user", "look")
+ atts = {
+ urow[0]: [
+ {
+ "attachment_id": "i1",
+ "kind": "image",
+ "mime_type": "image/png",
+ "filename": "x.png",
+ "content": _PNG,
+ }
+ ]
+ }
+ msgs = reconstruct_messages([urow], "ws1", atts)
+ assert msgs[0]["role"] == "user"
+ assert msgs[0]["content"][0] == {"type": "text", "text": "look"}
+ assert msgs[0]["content"][1]["type"] == "image_url"
+ assert msgs[0]["_attachments_meta"][0]["filename"] == "x.png"
+
+ def test_tool_row_multipart_image(self):
+ tc = json.dumps([{"id": "c1", "function": {"name": "read_file", "arguments": "{}"}}])
+ arow = _row("assistant", None, tool_calls=tc)
+ trow = _row("tool", "Image file: dog.png", tc_id="c1")
+ atts = {
+ trow[0]: [
+ {
+ "attachment_id": "i1",
+ "kind": "image",
+ "mime_type": "image/png",
+ "filename": "dog.png",
+ "content": _PNG,
+ }
+ ]
+ }
+ msgs = reconstruct_messages([arow, trow], "ws1", atts)
+ tool_msg = next(m for m in msgs if m["role"] == "tool")
+ assert isinstance(tool_msg["content"], list)
+ assert tool_msg["content"][0] == {"type": "text", "text": "Image file: dog.png"}
+ assert tool_msg["content"][1]["type"] == "image_url"
+ # Tool rows do NOT carry _attachments_meta (that's a user-display sibling).
+ assert "_attachments_meta" not in tool_msg
+
+ def test_tool_row_without_attachments_stays_string(self):
+ trow = _row("tool", "plain", tc_id="c1")
+ msgs = reconstruct_messages([trow], "ws1", None, repair=False)
+ assert msgs[0]["content"] == "plain"
diff --git a/tests/test_server_attachments_endpoints.py b/tests/test_server_attachments_endpoints.py
index 34aacf73..7a9fc698 100644
--- a/tests/test_server_attachments_endpoints.py
+++ b/tests/test_server_attachments_endpoints.py
@@ -10,7 +10,6 @@ from __future__ import annotations
import queue
import threading
-import uuid
from unittest.mock import MagicMock
import pytest
@@ -84,11 +83,18 @@ def app_client(tmp_path):
skip_permissions=False,
jwt_secret=_TEST_JWT_SECRET,
)
+ # Pending uploads live in the process-global per-node buffer now; clear it
+ # so staged uploads can't leak across tests.
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+
+ get_attachment_buffer()._entries.clear()
+
client = TestClient(app, raise_server_exceptions=False)
try:
yield client, mock_mgr
finally:
client.close()
+ get_attachment_buffer()._entries.clear()
reset_storage()
@@ -224,52 +230,10 @@ class TestUploadRejections:
assert resp.status_code == 200
-class TestPendingCap:
- def test_tenth_attachment_accepted_eleventh_rejected(self, app_client):
- client, _ = app_client
- for i in range(10):
- resp = client.post(
- "/v1/api/workstreams/ws-A/attachments",
- files={"file": (f"n{i}.md", b"x", "text/markdown")},
- headers=_auth("userA"),
- )
- assert resp.status_code == 200, resp.text
- resp = client.post(
- "/v1/api/workstreams/ws-A/attachments",
- files={"file": ("overflow.md", b"x", "text/markdown")},
- headers=_auth("userA"),
- )
- assert resp.status_code == 409
- assert resp.json().get("code") == "too_many"
-
- def test_cap_is_serialized_under_concurrent_uploads(self, app_client):
- # Pre-fill to cap-1, then fire two concurrent uploads. Exactly
- # one must succeed; the other must be rejected with 409.
- client, _ = app_client
- for i in range(9):
- assert (
- client.post(
- "/v1/api/workstreams/ws-A/attachments",
- files={"file": (f"pre{i}.md", b"x", "text/markdown")},
- headers=_auth("userA"),
- ).status_code
- == 200
- )
-
- import concurrent.futures
-
- def attempt(idx: int) -> int:
- return client.post(
- "/v1/api/workstreams/ws-A/attachments",
- files={"file": (f"race{idx}.md", b"x", "text/markdown")},
- headers=_auth("userA"),
- ).status_code
-
- with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
- futures = [ex.submit(attempt, i) for i in range(2)]
- results = sorted(f.result() for f in futures)
- # One success (200) + one cap-exceeded (409); never 200+200.
- assert results == [200, 409]
+# (The per-user pending-upload cap was removed with the content-addressing
+# cutover — pending uploads live in the per-node buffer, bounded by its own
+# size/TTL ceilings rather than a per-(ws,user) count. The cap tests that
+# lived here are gone.)
# ---------------------------------------------------------------------------
@@ -552,21 +516,28 @@ class TestSendMessageAttachments:
atts = captured["attachments"]
assert [x.attachment_id for x in atts] == [c, a, b]
- def test_send_oversized_attachment_ids_list_rejected(self, app_client):
- # Hostile / buggy clients should not be able to push an
- # arbitrarily long IN (...) clause through reservation.
+ def test_send_unknown_ids_resolve_to_nothing(self, app_client):
+ # The old oversized-IN-clause / cap rejection is gone (no DB
+ # reservation, no per-user cap). Unknown ids simply don't resolve
+ # from the buffer — the send proceeds with no attachments rather
+ # than 400-ing.
client, mgr = app_client
- self._wire_ws(mgr, "ws-A", "userA")
- from turnstone.core.attachments import MAX_PENDING_ATTACHMENTS_PER_USER_WS
-
- too_many = [f"id-{i}" for i in range(MAX_PENDING_ATTACHMENTS_PER_USER_WS + 1)]
+ captured, _ = self._wire_ws(mgr, "ws-A", "userA")
+ many = [f"id-{i}" for i in range(50)]
resp = client.post(
"/v1/api/workstreams/ws-A/send",
- json={"message": "x", "attachment_ids": too_many},
+ json={"message": "x", "attachment_ids": many},
headers=_auth("userA"),
)
- assert resp.status_code == 400
- assert resp.json().get("code") == "too_many"
+ assert resp.status_code == 200
+
+ import time
+
+ for _ in range(50):
+ if "attachments" in captured:
+ break
+ time.sleep(0.01)
+ assert captured["attachments"] is None
def test_send_forged_id_from_other_user_ignored(self, app_client):
client, mgr = app_client
@@ -666,13 +637,13 @@ class TestQueuedSendWithAttachments:
assert captured["attachment_ids"] == [b, a]
-class TestQueuedAttachmentReservation:
- """Once a queued send reserves its attachments, concurrent operations
- (delete, auto-consume, explicit reuse) must not disturb them."""
+class TestBusyWorkerAttachments:
+ """An attachment-bearing send to a busy worker can't ride the text-only
+ queue seam — it returns ``attachments_busy`` and the staged bytes stay in
+ the buffer (a peek, not a drain) so the client can retry once idle."""
def _wire_busy_ws(self, mgr, ws_id: str):
- """Mock ws whose worker is always alive (forces queue path).
- Uses the real ChatSession's queue_message so reservation runs."""
+ """Mock ws whose worker is always alive (forces the queue path)."""
from turnstone.core.session import ChatSession
from turnstone.core.workstream import WorkstreamState
@@ -686,8 +657,6 @@ class TestQueuedAttachmentReservation:
tool_timeout=10,
user_id="userA",
)
- # Pin the session to ws-A so queue_message / dequeue_message
- # operate on the expected storage rows.
session._ws_id = ws_id
ui = MagicMock()
@@ -708,119 +677,7 @@ class TestQueuedAttachmentReservation:
mgr.get.return_value = ws
return ws, session
- def _reserve_attachment(self, client, mgr, ws_id: str, filename: str = "q.md"):
- """Set up a reserved attachment for the busy-worker tests below.
-
- The queue-with-attachments path was removed (queued user turns
- can't carry attachments — see ``AttachmentsNotQueueableError``),
- so the tests reserve directly via ``reserve_attachments`` to
- produce the same on-disk state without going through the
- rejected route path.
- """
- from turnstone.core.memory import reserve_attachments
-
- aid = _upload(client, ws_id, "userA", filename, b"Q", "text/markdown")
- ws, session = self._wire_busy_ws(mgr, ws_id)
- msg_id = uuid.uuid4().hex
- reserve_attachments([aid], msg_id, ws_id, "userA")
- return aid, msg_id, session
-
- def test_reserved_attachment_hidden_from_pending_listing(self, app_client):
- client, mgr = app_client
- aid, _mid, _session = self._reserve_attachment(client, mgr, "ws-A")
- resp = client.get("/v1/api/workstreams/ws-A/attachments", headers=_auth("userA"))
- # Reserved attachment is not in the pending listing
- ids = [a["attachment_id"] for a in resp.json()["attachments"]]
- assert aid not in ids
-
- def test_reserved_attachment_cannot_be_deleted(self, app_client):
- client, mgr = app_client
- aid, _mid, _session = self._reserve_attachment(client, mgr, "ws-A")
- resp = client.delete(
- f"/v1/api/workstreams/ws-A/attachments/{aid}",
- headers=_auth("userA"),
- )
- # Delete silently masks reserved ones as not-found (can't delete
- # while tied to a queued message).
- assert resp.status_code == 404
- # Still exists on the backend
- from turnstone.core.memory import get_attachment
-
- assert get_attachment(aid) is not None
-
- def test_reserved_attachment_not_auto_consumed_by_later_send(self, app_client):
- client, mgr = app_client
- aid, _mid, session = self._reserve_attachment(client, mgr, "ws-A")
-
- # Swap the busy worker for an idle one and capture the next
- # session.send call so we can assert on its attachment list.
- captured: dict = {}
-
- def fake_send(message, attachments=None, send_id=None):
- captured["message"] = message
- captured["attachments"] = attachments
- captured["send_id"] = send_id
-
- session.send = fake_send # type: ignore[method-assign]
- ws = mgr.get.return_value
- ws.worker_thread = None # idle → non-queue path
- ws._worker_running = False
-
- # Auto-consume on a follow-up send: reserved attachment must not
- # be picked up (another turn isn't entitled to it).
- resp = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "follow up"},
- headers=_auth("userA"),
- )
- assert resp.status_code == 200
- import time
-
- for _ in range(50):
- if "attachments" in captured:
- break
- time.sleep(0.01)
- atts = captured.get("attachments")
- if atts is not None:
- assert aid not in [a.attachment_id for a in atts]
-
- def test_reserved_attachment_rejected_in_explicit_ids(self, app_client):
- client, mgr = app_client
- aid, _mid, session = self._reserve_attachment(client, mgr, "ws-A")
-
- captured: dict = {}
-
- def fake_send(message, attachments=None, send_id=None):
- captured["attachments"] = attachments
-
- session.send = fake_send # type: ignore[method-assign]
- ws = mgr.get.return_value
- ws.worker_thread = None
- ws._worker_running = False
-
- # A second send explicitly naming the reserved id: scope check
- # rejects it, so the attachment list is empty.
- resp = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "take mine", "attachment_ids": [aid]},
- headers=_auth("userA"),
- )
- assert resp.status_code == 200
- import time
-
- for _ in range(50):
- if "attachments" in captured:
- break
- time.sleep(0.01)
- atts = captured.get("attachments")
- # Either None (all scope-rejected → empty list collapses to None)
- # or an empty list — never contains the reserved id.
- if atts is not None:
- assert aid not in [a.attachment_id for a in atts]
-
def test_send_with_attachments_to_busy_worker_returns_attachments_busy(self, app_client):
- """An attempt to attach mid-tool-call returns ``attachments_busy``;
- attachments stay pending so the client can retry once idle."""
client, mgr = app_client
aid = _upload(client, "ws-A", "userA", "x.md", b"X", "text/markdown")
self._wire_busy_ws(mgr, "ws-A")
@@ -834,194 +691,13 @@ class TestQueuedAttachmentReservation:
assert body["status"] == "attachments_busy"
assert body["attached_ids"] == []
assert body["dropped_attachment_ids"] == [aid]
- # Reservation released — attachment is still pending and visible.
+ # The staged upload was peeked (not drained), so it's still pending
+ # and visible for a retry once the worker idles.
resp = client.get("/v1/api/workstreams/ws-A/attachments", headers=_auth("userA"))
ids = [a["attachment_id"] for a in resp.json()["attachments"]]
assert aid in ids
-class TestReserveThenDispatchRace:
- """Reservation happens BEFORE queue_message / worker start, so an
- overlapping request can't select the same row."""
-
- def test_overlapping_idle_send_cannot_resteal(self, app_client):
- # Kick off an idle send that reserves attachment A but blocks
- # inside session.send — then a second send with the same
- # explicit id must NOT receive A.
- client, mgr = app_client
- from turnstone.core.session import ChatSession
- from turnstone.core.workstream import WorkstreamState
-
- aid = _upload(client, "ws-A", "userA", "hold.md", b"hold", "text/markdown")
-
- # Real ChatSession for the idle path (reservation must be real)
- session = ChatSession(
- client=MagicMock(),
- model="test-model",
- ui=MagicMock(),
- instructions=None,
- temperature=0.3,
- max_tokens=1024,
- tool_timeout=10,
- user_id="userA",
- )
- session._ws_id = "ws-A"
-
- first_captured: dict = {}
- gate = threading.Event()
-
- def first_send(message, attachments=None, send_id=None):
- first_captured["attachments"] = attachments
- first_captured["send_id"] = send_id
- gate.wait(timeout=5.0) # Hold the worker so #2 races against us
-
- session.send = first_send # type: ignore[method-assign]
-
- ui = MagicMock()
- ui._ws_lock = threading.Lock()
- ui._ws_messages = 0
- ui._ws_turn_tool_calls = 0
-
- ws = MagicMock()
- ws.id = "ws-A"
- ws.state = WorkstreamState.IDLE
- ws.ui = ui
- ws.session = session
- ws.worker_thread = None
- ws._worker_running = False
- ws._lock = threading.RLock()
- mgr.get.return_value = ws
-
- # First send — reserves A under its send_id, worker blocks
- resp1 = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "one", "attachment_ids": [aid]},
- headers=_auth("userA"),
- )
- assert resp1.status_code == 200
-
- import time
-
- for _ in range(50):
- if first_captured.get("attachments") is not None:
- break
- time.sleep(0.01)
- assert first_captured["attachments"] is not None
- assert [a.attachment_id for a in first_captured["attachments"]] == [aid]
- first_send_id = first_captured["send_id"]
- assert first_send_id
-
- # Second send — idle path busy-check still sees the mock's
- # worker as "not alive" (we didn't update it) so this enters
- # the idle branch. Reservation must skip the already-reserved
- # row, so send sees no attachments.
- second_captured: dict = {}
-
- def second_send(message, attachments=None, send_id=None):
- second_captured["attachments"] = attachments
-
- session.send = second_send # type: ignore[method-assign]
-
- resp2 = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "two", "attachment_ids": [aid]},
- headers=_auth("userA"),
- )
- assert resp2.status_code == 200
-
- for _ in range(50):
- if "attachments" in second_captured:
- break
- time.sleep(0.01)
- # The second send either got no attachments or an empty list —
- # it never saw the row that the first send reserved.
- seen = second_captured.get("attachments")
- assert not seen or aid not in [a.attachment_id for a in (seen or [])]
-
- # Release the first worker so the fixture can tear down cleanly
- gate.set()
-
- def test_worker_exception_releases_reservation(self, app_client):
- # If session.send raises inside the worker thread, the
- # reservation must be released so the attachment isn't
- # permanently soft-locked.
- client, mgr = app_client
- from turnstone.core.memory import get_attachment
-
- aid = _upload(client, "ws-A", "userA", "boom.md", b"x", "text/markdown")
-
- captured, session = TestSendMessageAttachments._wire_ws(
- TestSendMessageAttachments, mgr, "ws-A", "userA"
- )
-
- def exploding_send(message, attachments=None, send_id=None):
- raise RuntimeError("worker blew up")
-
- session.send = exploding_send # type: ignore[method-assign]
-
- resp = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "boom", "attachment_ids": [aid]},
- headers=_auth("userA"),
- )
- assert resp.status_code == 200
-
- # Wait for the worker thread to finish (crash + cleanup).
- import time
-
- for _ in range(100):
- row = get_attachment(aid)
- if row and row.get("reserved_for_msg_id") is None:
- break
- time.sleep(0.01)
-
- row = get_attachment(aid)
- # Reservation must be released so the attachment is usable again.
- assert row is not None
- assert row["reserved_for_msg_id"] is None
- assert row["message_id"] is None
-
- def test_partial_reservation_proceeds_with_reserved_subset(self, app_client):
- # Pre-reserve one id; a send listing two explicit ids should
- # proceed with only the non-reserved one.
- client, mgr = app_client
- from turnstone.core.memory import reserve_attachments
-
- a = _upload(client, "ws-A", "userA", "a.md", b"A", "text/markdown")
- b = _upload(client, "ws-A", "userA", "b.md", b"B", "text/markdown")
- # Pre-reserve 'a' under a fake prior send
- reserve_attachments([a], "prior-send", "ws-A", "userA")
-
- captured: dict = {}
-
- def fake_send(message, attachments=None, send_id=None):
- captured["attachments"] = attachments
-
- ws_tuple = TestSendMessageAttachments._wire_ws(
- TestSendMessageAttachments, mgr, "ws-A", "userA"
- )
- captured = ws_tuple[0] # _wire_ws returns (captured, session)
- ws_tuple[1].send = fake_send # type: ignore[method-assign]
-
- resp = client.post(
- "/v1/api/workstreams/ws-A/send",
- json={"message": "both", "attachment_ids": [a, b]},
- headers=_auth("userA"),
- )
- assert resp.status_code == 200
-
- import time
-
- for _ in range(50):
- if "attachments" in captured:
- break
- time.sleep(0.01)
- atts = captured.get("attachments") or []
- ids = [x.attachment_id for x in atts]
- # Only the un-pre-reserved id survives
- assert ids == [b]
-
-
class TestServiceScopedActorFlow:
"""Service-scoped tokens bypass ownership checks and file attachments
under the workstream owner; send() must consume them using the same
@@ -1042,8 +718,9 @@ class TestServiceScopedActorFlow:
)
svc_headers = {"Authorization": f"Bearer {service_token}"}
- # Upload to ws-A (owned by userA) as service — file should land
- # under owner "userA", not "svc-bot"
+ # Upload to ws-A (owned by userA) as service — the staged upload is
+ # filed under the owner "userA" (the resolver's uid), not "svc-bot",
+ # so the later owner-resolved send can find it in the buffer.
resp = client.post(
"/v1/api/workstreams/ws-A/attachments",
files={"file": ("svc.md", b"svc", "text/markdown")},
@@ -1052,10 +729,11 @@ class TestServiceScopedActorFlow:
assert resp.status_code == 200
aid = resp.json()["attachment_id"]
- from turnstone.core.memory import get_attachment
+ from turnstone.core.attachment_buffer import get_attachment_buffer
- row = get_attachment(aid)
- assert row["user_id"] == "userA" # filed under the owner
+ # Staged under the owner uid (userA), not the service caller (svc-bot).
+ assert get_attachment_buffer().get(aid, ws_id="ws-A", user_id="userA") is not None
+ assert get_attachment_buffer().get(aid, ws_id="ws-A", user_id="svc-bot") is None
# Now drive /api/send as the service token — the resolver uses
# the ws owner (userA) to look up attachments, so the upload is
diff --git a/tests/test_server_attachments_on_create.py b/tests/test_server_attachments_on_create.py
index 18e0b66f..f11c8bfb 100644
--- a/tests/test_server_attachments_on_create.py
+++ b/tests/test_server_attachments_on_create.py
@@ -1,8 +1,8 @@
"""Tests for the multipart variant of POST /v1/api/workstreams/new.
Exercises:
-- The pure helpers `_validate_and_save_uploaded_files` and
- `_reserve_and_resolve_attachments` (added alongside the multipart path).
+- The pure helpers ``validate_and_save_uploaded_files`` (stages to the
+ per-node buffer) and ``resolve_staged_attachments`` (peeks them back).
- The full create endpoint via TestClient with a FakeSession factory so
the initial-message dispatch thread runs end-to-end without an LLM.
"""
@@ -54,15 +54,16 @@ def _auth(user: str) -> dict[str, str]:
class TestValidateAndSaveUploadedFiles:
- def test_saves_image_and_text(self, tmp_path):
+ def test_stages_image_and_text_to_buffer(self, tmp_path):
+ import hashlib
+
+ from turnstone.core.attachment_buffer import get_attachment_buffer
from turnstone.core.attachments import (
validate_and_save_uploaded_files as _validate_and_save_uploaded_files,
)
- from turnstone.core.memory import list_pending_attachments
- from turnstone.core.storage import init_storage, reset_storage
- reset_storage()
- init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
+ buf = get_attachment_buffer()
+ buf._entries.clear()
try:
files = [
("hi.png", "image/png", PNG_1x1),
@@ -71,12 +72,13 @@ class TestValidateAndSaveUploadedFiles:
ids, err = _validate_and_save_uploaded_files(files, "ws-X", "userA")
assert err is None
assert len(ids) == 2
- pending = list_pending_attachments("ws-X", "userA")
- assert len(pending) == 2
- kinds = {p["kind"] for p in pending}
- assert kinds == {"image", "text"}
+ # Ids are content hashes (content-addressed staging).
+ assert ids[0] == hashlib.sha256(PNG_1x1).hexdigest()
+ staged = buf.list_for(ws_id="ws-X", user_id="userA")
+ assert len(staged) == 2
+ assert {s.kind for s in staged} == {"image", "text"}
finally:
- reset_storage()
+ buf._entries.clear()
def test_rejects_oversized_image(self, tmp_path):
from turnstone.core.attachments import IMAGE_SIZE_CAP
@@ -116,76 +118,85 @@ class TestValidateAndSaveUploadedFiles:
finally:
reset_storage()
- def test_pending_cap_returns_409(self, tmp_path):
- from turnstone.core.attachments import MAX_PENDING_ATTACHMENTS_PER_USER_WS
- from turnstone.core.attachments import (
- validate_and_save_uploaded_files as _validate_and_save_uploaded_files,
+
+# (The per-user pending-upload cap was removed with the content-addressing
+# cutover; ``validate_and_save_uploaded_files`` no longer 409s — the buffer's
+# own size/TTL ceilings bound a flood.)
+
+
+class TestResolveStagedAttachments:
+ def _stage(self, ws_id, user_id, filename, mime, kind, content):
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+
+ return get_attachment_buffer().stage(
+ ws_id=ws_id,
+ user_id=user_id,
+ filename=filename,
+ mime_type=mime,
+ kind=kind,
+ content=content,
)
- from turnstone.core.memory import save_attachment
- from turnstone.core.storage import init_storage, reset_storage
- reset_storage()
- init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
- try:
- # Saturate the pending cap
- for i in range(MAX_PENDING_ATTACHMENTS_PER_USER_WS):
- save_attachment(
- f"pre-{i}", "ws-X", "userA", f"f{i}.txt", "text/plain", 1, "text", b"x"
- )
- files = [("notes.md", "text/markdown", b"hello")]
- ids, err = _validate_and_save_uploaded_files(files, "ws-X", "userA")
- assert err is not None
- assert err.status_code == 409
- assert ids == []
- finally:
- reset_storage()
-
-
-class TestReserveAndResolveAttachments:
- def test_reserves_and_returns_attachments(self, tmp_path):
+ def test_resolves_staged_to_attachments(self):
+ from turnstone.core.attachment_buffer import get_attachment_buffer
from turnstone.core.attachments import Attachment
from turnstone.core.attachments import (
- reserve_and_resolve_attachments as _reserve_and_resolve_attachments,
+ resolve_staged_attachments as _resolve_staged_attachments,
)
- from turnstone.core.memory import save_attachment
- from turnstone.core.storage import init_storage, reset_storage
- reset_storage()
- init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
+ buf = get_attachment_buffer()
+ buf._entries.clear()
try:
- save_attachment("a1", "ws-X", "userA", "a.txt", "text/plain", 5, "text", b"hello")
- save_attachment("a2", "ws-X", "userA", "b.png", "image/png", 91, "image", PNG_1x1)
- resolved, ordered, dropped = _reserve_and_resolve_attachments(
- ["a1", "a2"], "send-1", "ws-X", "userA"
+ a1 = self._stage("ws-X", "userA", "a.txt", "text/plain", "text", b"hello")
+ a2 = self._stage("ws-X", "userA", "b.png", "image/png", "image", PNG_1x1)
+ resolved, taken, dropped = _resolve_staged_attachments(
+ [a1.attachment_id, a2.attachment_id], "ws-X", "userA"
)
- assert ordered == ["a1", "a2"]
+ assert taken == [a1.attachment_id, a2.attachment_id]
assert dropped == []
assert len(resolved) == 2
assert all(isinstance(a, Attachment) for a in resolved)
- kinds = [a.kind for a in resolved]
- assert kinds == ["text", "image"]
+ assert [a.kind for a in resolved] == ["text", "image"]
finally:
- reset_storage()
+ buf._entries.clear()
- def test_double_reserve_drops_second(self, tmp_path):
+ def test_unknown_id_is_dropped(self):
+ from turnstone.core.attachment_buffer import get_attachment_buffer
from turnstone.core.attachments import (
- reserve_and_resolve_attachments as _reserve_and_resolve_attachments,
+ resolve_staged_attachments as _resolve_staged_attachments,
)
- from turnstone.core.memory import save_attachment
- from turnstone.core.storage import init_storage, reset_storage
- reset_storage()
- init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
+ buf = get_attachment_buffer()
+ buf._entries.clear()
try:
- save_attachment("a1", "ws-X", "userA", "a.txt", "text/plain", 5, "text", b"hello")
- r1, ord1, _ = _reserve_and_resolve_attachments(["a1"], "send-A", "ws-X", "userA")
- assert len(r1) == 1
- r2, ord2, drop2 = _reserve_and_resolve_attachments(["a1"], "send-B", "ws-X", "userA")
- assert r2 == []
- assert ord2 == []
- assert drop2 == ["a1"]
+ a1 = self._stage("ws-X", "userA", "a.txt", "text/plain", "text", b"hello")
+ resolved, taken, dropped = _resolve_staged_attachments(
+ [a1.attachment_id, "never-staged"], "ws-X", "userA"
+ )
+ assert taken == [a1.attachment_id]
+ assert dropped == ["never-staged"]
+ assert len(resolved) == 1
finally:
- reset_storage()
+ buf._entries.clear()
+
+ def test_cross_user_id_not_resolved(self):
+ # A staged id scoped to another user (same ws) must not resolve.
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.attachments import (
+ resolve_staged_attachments as _resolve_staged_attachments,
+ )
+
+ buf = get_attachment_buffer()
+ buf._entries.clear()
+ try:
+ other = self._stage("ws-X", "userB", "a.txt", "text/plain", "text", b"hello")
+ resolved, taken, dropped = _resolve_staged_attachments(
+ [other.attachment_id], "ws-X", "userA"
+ )
+ assert resolved == []
+ assert dropped == [other.attachment_id]
+ finally:
+ buf._entries.clear()
# ---------------------------------------------------------------------------
@@ -216,22 +227,36 @@ class _FakeSession:
def send(self, text, attachments=None, send_id=None):
with self._lock:
self.sends.append((text, list(attachments or []), send_id))
- # Simulate the real ChatSession's consume step against storage
- # so callers can assert the lifecycle landed.
- if attachments and send_id and self.ws_id and self.user_id:
- import uuid as _uuid
-
- from turnstone.core.memory import mark_attachments_consumed
-
- ids = [a.attachment_id for a in attachments]
- mark_attachments_consumed(
- ids,
- _uuid.uuid4().hex, # synthetic conversation message id
- self.ws_id,
- self.user_id,
- reserved_for_msg_id=send_id,
+ # Simulate the real ChatSession commit: write each attachment
+ # content-addressed, record the ref-list on a (synthetic) message
+ # row, and drain the staged handles from the buffer — so callers
+ # can assert the lifecycle landed.
+ if attachments and self.ws_id and self.user_id:
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.memory import (
+ save_attachment,
+ save_message,
+ set_message_attachments,
)
+ mid = save_message(self.ws_id, "user", text)
+ buf = get_attachment_buffer()
+ ref_ids = []
+ for a in attachments:
+ save_attachment(
+ a.attachment_id,
+ self.ws_id,
+ self.user_id,
+ a.filename,
+ a.mime_type,
+ len(a.content),
+ a.kind,
+ a.content,
+ )
+ ref_ids.append(a.attachment_id)
+ buf.discard(a.attachment_id, ws_id=self.ws_id, user_id=self.user_id)
+ set_message_attachments(self.ws_id, mid, ref_ids)
+
# Methods the create handler may call but we don't care about
def set_watch_runner(self, *_a, **_kw):
pass
@@ -318,17 +343,27 @@ def app_client(tmp_path, monkeypatch):
skip_permissions=False,
jwt_secret=_TEST_JWT_SECRET,
)
+ # Pending uploads live in the process-global per-node buffer; clear it so
+ # staged uploads can't leak across tests.
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+
+ get_attachment_buffer()._entries.clear()
+
client = TestClient(app, raise_server_exceptions=False)
try:
yield client, fake_sessions, gq
finally:
client.close()
+ get_attachment_buffer()._entries.clear()
reset_storage()
class TestCreateMultipart:
def test_create_with_image_and_initial_message(self, app_client):
- from turnstone.core.memory import list_pending_attachments
+ import hashlib
+
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.memory import attachment_referenced_in_ws, get_attachment
client, sessions, _gq = app_client
meta = {"name": "demo", "initial_message": "describe this image"}
@@ -343,6 +378,8 @@ class TestCreateMultipart:
ws_id = data["ws_id"]
assert ws_id
assert len(data["attachment_ids"]) == 1
+ aid = data["attachment_ids"][0]
+ assert aid == hashlib.sha256(PNG_1x1).hexdigest() # content-addressed id
# Wait briefly for the dispatch thread
deadline = time.time() + 2.0
@@ -355,16 +392,27 @@ class TestCreateMultipart:
assert sessions[0].sends, "session.send was not invoked"
text, atts, send_id = sessions[0].sends[0]
assert text == "describe this image"
- assert send_id # reservation token threaded through
+ assert send_id # tracking token threaded through
assert len(atts) == 1
assert atts[0].kind == "image"
- # Lifecycle: the FakeSession marks them consumed via storage —
- # so the pending-list for this ws should be empty after dispatch.
- assert list_pending_attachments(ws_id, "userA") == []
+ # Lifecycle: the FakeSession commit wrote the blob content-addressed +
+ # recorded the ref-list + drained the buffer. Poll for the worker.
+ deadline = time.time() + 2.0
+ while time.time() < deadline:
+ if get_attachment(aid) is not None and attachment_referenced_in_ws(aid, ws_id):
+ break
+ time.sleep(0.02)
+ row = get_attachment(aid)
+ assert row is not None and row["refcount"] >= 1
+ assert attachment_referenced_in_ws(aid, ws_id) is True
+ # Drained from the buffer post-commit.
+ assert get_attachment_buffer().get(aid, ws_id=ws_id, user_id="userA") is None
- def test_create_with_attachments_no_initial_message_keeps_pending(self, app_client):
- from turnstone.core.memory import list_pending_attachments
+ def test_create_with_attachments_no_initial_message_keeps_staged(self, app_client):
+ import hashlib
+
+ from turnstone.core.attachment_buffer import get_attachment_buffer
client, _, _gq = app_client
meta = {"name": "stash"}
@@ -377,9 +425,11 @@ class TestCreateMultipart:
assert resp.status_code == 200, resp.text
data = resp.json()
ws_id = data["ws_id"]
- pending = list_pending_attachments(ws_id, "userA")
- assert len(pending) == 1
- assert pending[0]["filename"] == "notes.md"
+ assert data["attachment_ids"] == [hashlib.sha256(b"# hello\n").hexdigest()]
+ # No initial_message → no send → the upload stays staged in the buffer.
+ staged = get_attachment_buffer().list_for(ws_id=ws_id, user_id="userA")
+ assert len(staged) == 1
+ assert staged[0].filename == "notes.md"
def test_create_rejects_oversized_image_and_rolls_back(self, app_client):
from turnstone.core.attachments import IMAGE_SIZE_CAP
diff --git a/tests/test_session_attachments.py b/tests/test_session_attachments.py
index 652567c4..e4f3c6e8 100644
--- a/tests/test_session_attachments.py
+++ b/tests/test_session_attachments.py
@@ -9,9 +9,7 @@ import pytest
from turnstone.core.attachments import Attachment
from turnstone.core.memory import (
get_attachment,
- list_pending_attachments,
register_workstream,
- save_attachment,
)
from turnstone.core.session import ChatSession
@@ -132,23 +130,19 @@ class TestMultipartBuild:
class TestPersistenceAndConsumption:
- def test_db_row_stores_text_only(self, tmp_db, mock_openai_client):
+ def test_db_row_stores_text_only_and_records_ref_list(self, tmp_db, mock_openai_client):
+ import hashlib
+
s = _make_session(mock_openai_client)
- save_attachment(
- "att-persist",
- s._ws_id,
- "u1",
- "note.md",
- "text/markdown",
- 5,
- "text",
- b"hello",
- )
- att = Attachment("att-persist", "note.md", "text/markdown", "text", b"hello")
+ content = b"hello"
+ aid = hashlib.sha256(content).hexdigest() # the content hash is the id
+ att = Attachment(aid, "note.md", "text/markdown", "text", content)
_run_send(s, "user text", attachments=[att])
- # The conversations row's text content is just the user input —
- # the attachment is linked separately via message_id.
+ # The conversations row's text content is just the user input; the
+ # attachment is linked via the ``attachments`` ref-list column.
+ import json
+
import sqlalchemy as sa
from turnstone.core.storage._registry import get_storage
@@ -156,41 +150,60 @@ class TestPersistenceAndConsumption:
with get_storage()._conn() as conn:
rows = conn.execute(
- sa.select(conversations.c.content, conversations.c.id)
+ sa.select(conversations.c.content, conversations.c.id, conversations.c.attachments)
.where(conversations.c.ws_id == s._ws_id)
.order_by(conversations.c.id)
).fetchall()
assert len(rows) == 1
assert rows[0][0] == "user text"
- msg_id = rows[0][1]
+ assert json.loads(rows[0][2]) == [aid]
- # Attachment should be consumed and linked to the message
- assert list_pending_attachments(s._ws_id, "u1") == []
- att_row = get_attachment("att-persist")
+ # The blob was written content-addressed at refcount 1, origin upload.
+ att_row = get_attachment(aid)
assert att_row is not None
- assert att_row["message_id"] == msg_id
+ assert att_row["content"] == content
+ assert att_row["refcount"] == 1
+ assert att_row["origin"] == "upload"
- def test_consumption_scoped_to_user(self, tmp_db, mock_openai_client):
- # A session running as user B must not consume user A's attachments
- # even if the id is in the list passed to send().
- s = _make_session(mock_openai_client, user_id="userB")
- save_attachment(
- "att-other",
- s._ws_id,
- "userA",
- "a.md",
- "text/plain",
- 1,
- "text",
- b"A",
+ def test_send_drains_the_upload_buffer(self, tmp_db, mock_openai_client):
+ # Bytes staged in the per-node buffer are drained (discarded) once the
+ # send commits them content-addressed — they don't linger as pending.
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+
+ s = _make_session(mock_openai_client)
+ buf = get_attachment_buffer()
+ staged = buf.stage(
+ ws_id=s._ws_id,
+ user_id=s._user_id,
+ filename="note.md",
+ mime_type="text/markdown",
+ kind="text",
+ content=b"buffered",
)
- # Session constructs multipart content regardless (trust-but-verify),
- # but the DB-level mark is scoped — attachment stays pending for A.
- att = Attachment("att-other", "a.md", "text/plain", "text", b"A")
- _run_send(s, "hi", attachments=[att])
- att_row = get_attachment("att-other")
- assert att_row is not None
- assert att_row["message_id"] is None
+ assert buf.get(staged.attachment_id, ws_id=s._ws_id, user_id=s._user_id) is not None
+ att = Attachment(staged.attachment_id, "note.md", "text/markdown", "text", b"buffered")
+ _run_send(s, "user text", attachments=[att])
+ # Drained from the buffer post-commit.
+ assert buf.get(staged.attachment_id, ws_id=s._ws_id, user_id=s._user_id) is None
+
+ def test_reload_reconstructs_multipart(self, tmp_db, mock_openai_client):
+ import hashlib
+
+ from turnstone.core.memory import load_messages
+
+ s = _make_session(mock_openai_client)
+ content = b"# doc\n"
+ aid = hashlib.sha256(content).hexdigest()
+ att = Attachment(aid, "d.md", "text/markdown", "text", content)
+ _run_send(s, "see doc", attachments=[att])
+
+ msgs = load_messages(s._ws_id, repair=False)
+ assert msgs[0]["role"] == "user"
+ parts = msgs[0]["content"]
+ assert isinstance(parts, list)
+ assert parts[0] == {"type": "text", "text": "see doc"}
+ assert parts[1]["type"] == "document"
+ assert parts[1]["document"]["data"] == "# doc\n"
class TestProviderIntegration:
@@ -286,7 +299,8 @@ class TestQueuedAttachmentsRejected:
from turnstone.core.session import AttachmentsNotQueueableError
s = _make_session(mock_openai_client)
- save_attachment("a-q1", s._ws_id, "u1", "q.md", "text/markdown", 1, "text", b"q")
+ # Rejection is on the ``attachment_ids`` argument alone — no row need
+ # exist (the buffer is the pending store; queueing never touches it).
with pytest.raises(AttachmentsNotQueueableError):
s.queue_message("queued text", attachment_ids=["a-q1"])
# Queue stayed empty — nothing partially committed.
diff --git a/tests/test_storage_attachments.py b/tests/test_storage_attachments.py
index bc20a494..27315591 100644
--- a/tests/test_storage_attachments.py
+++ b/tests/test_storage_attachments.py
@@ -1,16 +1,19 @@
-"""Tests for workstream_attachments storage layer."""
+"""Tests for the content-addressed, refcounted workstream_attachments store.
+
+The pre-cutover persisted pending/reserved/consumed lifecycle (message_id /
+reserved_* + the per-user upload cap) is gone — pending uploads now live in the
+per-node in-memory buffer (see ``test_attachment_buffer.py``), and storage holds
+only committed blobs: written content-addressed at send-commit, deduped by
+content hash, and reference-counted via the ``conversations.attachments``
+ref-list. These tests pin that model at the storage boundary.
+"""
from __future__ import annotations
-import uuid
+import hashlib
import pytest
-
-def _aid() -> str:
- return uuid.uuid4().hex
-
-
PNG_1x1 = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\x9cc\xfc\xcf"
@@ -18,6 +21,11 @@ PNG_1x1 = (
)
+def _hash(content: bytes) -> str:
+ """The content-addressed id: bytes' sha256 hex (what the buffer computes)."""
+ return hashlib.sha256(content).hexdigest()
+
+
class TestSaveMessageReturnsId:
def test_returns_autoincrement_id(self, backend):
backend.register_workstream("ws-ret")
@@ -29,153 +37,142 @@ class TestSaveMessageReturnsId:
assert m2 > m1
-class TestAttachmentCRUD:
- def test_save_then_list_pending(self, backend):
- backend.register_workstream("ws-a")
- aid = _aid()
- backend.save_attachment(
- aid, "ws-a", "user-1", "hello.txt", "text/plain", 5, "text", b"hello"
- )
- pending = backend.list_pending_attachments("ws-a", "user-1")
- assert len(pending) == 1
- row = pending[0]
+class TestContentAddressedWrite:
+ def test_save_writes_blob_at_refcount_one(self, backend):
+ backend.register_workstream("ws-ca")
+ aid = _hash(b"hello")
+ backend.save_attachment(aid, "ws-ca", "u", "hello.txt", "text/plain", 5, "text", b"hello")
+ row = backend.get_attachment(aid)
+ assert row is not None
assert row["attachment_id"] == aid
- assert row["filename"] == "hello.txt"
- assert row["mime_type"] == "text/plain"
- assert row["size_bytes"] == 5
- assert row["kind"] == "text"
- # bytes must not leak into the pending-listing payload
- assert "content" not in row
+ assert row["content"] == b"hello"
+ assert row["refcount"] == 1
+ assert row["origin"] == "upload"
- def test_list_pending_isolates_users(self, backend):
- backend.register_workstream("ws-iso")
- a1 = _aid()
- a2 = _aid()
- backend.save_attachment(a1, "ws-iso", "user-A", "a.txt", "text/plain", 1, "text", b"A")
- backend.save_attachment(a2, "ws-iso", "user-B", "b.txt", "text/plain", 1, "text", b"B")
- a_pending = backend.list_pending_attachments("ws-iso", "user-A")
- b_pending = backend.list_pending_attachments("ws-iso", "user-B")
- assert [r["attachment_id"] for r in a_pending] == [a1]
- assert [r["attachment_id"] for r in b_pending] == [a2]
+ def test_origin_tool_recorded(self, backend):
+ backend.register_workstream("ws-origin")
+ aid = _hash(PNG_1x1)
+ backend.save_attachment(
+ aid, "ws-origin", "u", "t.png", "image/png", len(PNG_1x1), "image", PNG_1x1, "tool"
+ )
+ row = backend.get_attachment(aid)
+ assert row is not None
+ assert row["origin"] == "tool"
+ assert row["refcount"] == 1
- def test_get_attachments_bulk_returns_bytes(self, backend):
+ def test_identical_bytes_dedup_and_bump_refcount(self, backend):
+ backend.register_workstream("ws-dedup")
+ aid = _hash(b"same")
+ backend.save_attachment(aid, "ws-dedup", "u", "a.txt", "text/plain", 4, "text", b"same")
+ # A second reference to identical bytes does not duplicate the row —
+ # it bumps the refcount (e.g. two messages reference the same blob).
+ backend.save_attachment(aid, "ws-dedup", "u", "b.txt", "text/plain", 4, "text", b"same")
+ rows = backend.get_attachments([aid])
+ assert len(rows) == 1
+ assert rows[0]["refcount"] == 2
+ # First-writer metadata wins (INSERT-OR-IGNORE on the blob).
+ assert rows[0]["filename"] == "a.txt"
+
+ def test_distinct_bytes_are_distinct_blobs(self, backend):
+ backend.register_workstream("ws-distinct")
+ a1 = _hash(b"one")
+ a2 = _hash(b"two")
+ backend.save_attachment(a1, "ws-distinct", "u", "1.txt", "text/plain", 3, "text", b"one")
+ backend.save_attachment(a2, "ws-distinct", "u", "2.txt", "text/plain", 3, "text", b"two")
+ assert a1 != a2
+ rows = {r["attachment_id"]: r for r in backend.get_attachments([a1, a2])}
+ assert rows[a1]["content"] == b"one"
+ assert rows[a2]["content"] == b"two"
+
+
+class TestGetAttachments:
+ def test_bulk_returns_bytes(self, backend):
backend.register_workstream("ws-b")
- a1 = _aid()
- a2 = _aid()
+ a1 = _hash(b"one")
+ a2 = _hash(PNG_1x1)
backend.save_attachment(a1, "ws-b", "u", "one.txt", "text/plain", 3, "text", b"one")
backend.save_attachment(
a2, "ws-b", "u", "img.png", "image/png", len(PNG_1x1), "image", PNG_1x1
)
- rows = backend.get_attachments([a1, a2])
- by_id = {r["attachment_id"]: r for r in rows}
+ by_id = {r["attachment_id"]: r for r in backend.get_attachments([a1, a2])}
assert by_id[a1]["content"] == b"one"
assert by_id[a2]["content"] == PNG_1x1
assert by_id[a2]["kind"] == "image"
- def test_get_attachments_empty_input(self, backend):
+ def test_empty_input(self, backend):
assert backend.get_attachments([]) == []
+ def test_mixed_known_and_unknown_ids(self, backend):
+ backend.register_workstream("ws-mix")
+ known = _hash(b"k")
+ backend.save_attachment(known, "ws-mix", "u", "k.txt", "text/plain", 1, "text", b"k")
+ rows = backend.get_attachments([known, _hash(b"nope"), "definitely-not-an-id"])
+ assert len(rows) == 1
+ assert rows[0]["attachment_id"] == known
+
def test_get_attachment_missing_returns_none(self, backend):
assert backend.get_attachment("no-such-id") is None
- def test_delete_pending(self, backend):
- backend.register_workstream("ws-d")
- aid = _aid()
- backend.save_attachment(aid, "ws-d", "u", "x.txt", "text/plain", 1, "text", b"x")
- assert backend.delete_attachment(aid, "ws-d", "u") is True
- assert backend.list_pending_attachments("ws-d", "u") == []
- def test_delete_wrong_user_is_noop(self, backend):
- backend.register_workstream("ws-perm")
- aid = _aid()
- backend.save_attachment(aid, "ws-perm", "owner", "o.txt", "text/plain", 1, "text", b"o")
- assert backend.delete_attachment(aid, "ws-perm", "intruder") is False
- assert len(backend.list_pending_attachments("ws-perm", "owner")) == 1
+class TestSetMessageAttachments:
+ def test_records_ordered_ref_list(self, backend):
+ import json
- def test_delete_after_consumed_is_noop(self, backend):
- backend.register_workstream("ws-con")
- aid = _aid()
- backend.save_attachment(aid, "ws-con", "u", "c.txt", "text/plain", 1, "text", b"c")
- msg_id = backend.save_message("ws-con", "user", "hi")
- backend.mark_attachments_consumed([aid], msg_id, "ws-con", "u")
- assert backend.delete_attachment(aid, "ws-con", "u") is False
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["message_id"] == msg_id
+ import sqlalchemy as sa
+ from turnstone.core.storage._schema import conversations
-class TestConsumptionLinkage:
- def test_mark_consumed_links_message(self, backend):
- backend.register_workstream("ws-link")
- aid = _aid()
- backend.save_attachment(aid, "ws-link", "u", "f.txt", "text/plain", 1, "text", b"f")
- msg_id = backend.save_message("ws-link", "user", "with attach")
- backend.mark_attachments_consumed([aid], msg_id, "ws-link", "u")
+ backend.register_workstream("ws-ref")
+ mid = backend.save_message("ws-ref", "user", "hi")
+ a1, a2 = _hash(b"x"), _hash(b"y")
+ backend.save_attachment(a1, "ws-ref", "u", "x.txt", "text/plain", 1, "text", b"x")
+ backend.save_attachment(a2, "ws-ref", "u", "y.txt", "text/plain", 1, "text", b"y")
+ backend.set_message_attachments("ws-ref", mid, [a2, a1]) # order matters
+ with backend._conn() as conn:
+ raw = conn.execute(
+ sa.select(conversations.c.attachments).where(conversations.c.id == mid)
+ ).scalar_one()
+ assert json.loads(raw) == [a2, a1]
- # No longer listed as pending
- assert backend.list_pending_attachments("ws-link", "u") == []
- # Second mark is a no-op (won't re-link to a different message)
- other_msg_id = backend.save_message("ws-link", "user", "another")
- backend.mark_attachments_consumed([aid], other_msg_id, "ws-link", "u")
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["message_id"] == msg_id
+ def test_empty_input_is_noop(self, backend):
+ backend.register_workstream("ws-ref2")
+ mid = backend.save_message("ws-ref2", "user", "hi")
+ backend.set_message_attachments("ws-ref2", mid, []) # must not raise
+ # Column stays NULL → load yields a plain string message.
+ assert backend.load_messages("ws-ref2")[0]["content"] == "hi"
- def test_mark_consumed_empty_input(self, backend):
- backend.mark_attachments_consumed([], 0, "ws", "u") # must not raise
+ def test_scoped_to_ws(self, backend):
+ # A cross-ws message id is not written (defense-in-depth).
+ import sqlalchemy as sa
- def test_mark_consumed_wrong_user_is_noop(self, backend):
- backend.register_workstream("ws-scope")
- aid = _aid()
- backend.save_attachment(aid, "ws-scope", "owner", "o.txt", "text/plain", 1, "text", b"o")
- msg_id = backend.save_message("ws-scope", "user", "hi")
- # Different user tries to consume — must not link
- backend.mark_attachments_consumed([aid], msg_id, "ws-scope", "intruder")
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["message_id"] is None
+ from turnstone.core.storage._schema import conversations
- def test_mark_consumed_wrong_ws_is_noop(self, backend):
- backend.register_workstream("ws-scope2")
- backend.register_workstream("ws-other")
- aid = _aid()
- backend.save_attachment(aid, "ws-scope2", "u", "x.txt", "text/plain", 1, "text", b"x")
- msg_id = backend.save_message("ws-other", "user", "hi")
- # Try to link to a message in a different ws — must not succeed
- backend.mark_attachments_consumed([aid], msg_id, "ws-other", "u")
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["message_id"] is None
+ backend.register_workstream("ws-a")
+ backend.register_workstream("ws-b")
+ mid = backend.save_message("ws-a", "user", "hi")
+ aid = _hash(b"x")
+ backend.save_attachment(aid, "ws-a", "u", "x.txt", "text/plain", 1, "text", b"x")
+ backend.set_message_attachments("ws-b", mid, [aid]) # wrong ws
+ with backend._conn() as conn:
+ raw = conn.execute(
+ sa.select(conversations.c.attachments).where(conversations.c.id == mid)
+ ).scalar_one()
+ assert raw is None
class TestLoadMessagesReconstructsMultipart:
def test_user_message_with_image_and_text_doc(self, backend):
backend.register_workstream("ws-multi")
msg_id = backend.save_message("ws-multi", "user", "look at these")
-
- img_id = _aid()
- doc_id = _aid()
+ img_id = _hash(PNG_1x1)
+ doc_id = _hash(b"# hi\n")
backend.save_attachment(
- img_id,
- "ws-multi",
- "u",
- "tiny.png",
- "image/png",
- len(PNG_1x1),
- "image",
- PNG_1x1,
+ img_id, "ws-multi", "u", "tiny.png", "image/png", len(PNG_1x1), "image", PNG_1x1
)
backend.save_attachment(
- doc_id,
- "ws-multi",
- "u",
- "notes.md",
- "text/markdown",
- 5,
- "text",
- b"# hi\n",
+ doc_id, "ws-multi", "u", "notes.md", "text/markdown", 5, "text", b"# hi\n"
)
- backend.mark_attachments_consumed([img_id, doc_id], msg_id, "ws-multi", "u")
+ backend.set_message_attachments("ws-multi", msg_id, [img_id, doc_id])
msgs = backend.load_messages("ws-multi")
assert len(msgs) == 1
@@ -184,7 +181,6 @@ class TestLoadMessagesReconstructsMultipart:
content = user_msg["content"]
assert isinstance(content, list)
assert content[0] == {"type": "text", "text": "look at these"}
- # Image part: base64 data URI
kinds = [p["type"] for p in content[1:]]
assert "image_url" in kinds
assert "document" in kinds
@@ -195,198 +191,227 @@ class TestLoadMessagesReconstructsMultipart:
assert doc_part["document"]["media_type"] == "text/markdown"
assert doc_part["document"]["data"] == "# hi\n"
+ def test_ref_list_order_preserved(self, backend):
+ backend.register_workstream("ws-order")
+ mid = backend.save_message("ws-order", "user", "ordered")
+ a, b = _hash(b"AAA"), _hash(b"BBB")
+ backend.save_attachment(a, "ws-order", "u", "a.md", "text/markdown", 3, "text", b"AAA")
+ backend.save_attachment(b, "ws-order", "u", "b.md", "text/markdown", 3, "text", b"BBB")
+ # Record b before a — reconstruction must follow the ref-list order.
+ backend.set_message_attachments("ws-order", mid, [b, a])
+ docs = [
+ p for p in backend.load_messages("ws-order")[0]["content"] if p["type"] == "document"
+ ]
+ assert [d["document"]["data"] for d in docs] == ["BBB", "AAA"]
+
def test_user_message_without_attachments_stays_string(self, backend):
backend.register_workstream("ws-plain")
backend.save_message("ws-plain", "user", "plain text")
- msgs = backend.load_messages("ws-plain")
- assert msgs[0]["content"] == "plain text"
+ assert backend.load_messages("ws-plain")[0]["content"] == "plain text"
def test_invalid_utf8_text_attachment_shows_placeholder(self, backend):
backend.register_workstream("ws-bad")
msg_id = backend.save_message("ws-bad", "user", "oops")
- aid = _aid()
+ aid = _hash(b"\xff\xfe")
backend.save_attachment(aid, "ws-bad", "u", "bad.txt", "text/plain", 2, "text", b"\xff\xfe")
- backend.mark_attachments_consumed([aid], msg_id, "ws-bad", "u")
- msgs = backend.load_messages("ws-bad")
- # Undecodable text → placeholder so the user sees the attachment existed
- content = msgs[0]["content"]
+ backend.set_message_attachments("ws-bad", msg_id, [aid])
+ content = backend.load_messages("ws-bad")[0]["content"]
assert isinstance(content, list)
assert content[0] == {"type": "text", "text": "oops"}
assert content[1] == {"type": "text", "text": "[unreadable attachment: bad.txt]"}
+ def test_missing_blob_is_skipped(self, backend):
+ # A ref-list id whose blob was pruned (refcount hit 0 via another
+ # message's GC) reconstructs as plain text, not a crash.
+ backend.register_workstream("ws-missing")
+ mid = backend.save_message("ws-missing", "user", "gone")
+ backend.set_message_attachments("ws-missing", mid, [_hash(b"never-written")])
+ assert backend.load_messages("ws-missing")[0]["content"] == "gone"
-class TestDeleteWorkstreamCascade:
- def test_attachments_removed_on_workstream_delete(self, backend):
- backend.register_workstream("ws-cas")
- aid = _aid()
- backend.save_attachment(aid, "ws-cas", "u", "a.txt", "text/plain", 1, "text", b"a")
- msg_id = backend.save_message("ws-cas", "user", "hi")
- backend.mark_attachments_consumed([aid], msg_id, "ws-cas", "u")
- assert backend.delete_workstream("ws-cas") is True
- assert backend.get_attachment(aid) is None
+class TestToolImageReconstruction:
+ def test_tool_row_with_image_rebuilds_multipart(self, backend):
+ """Tool vision output is persisted content-addressed + referenced on the
+ tool row, so a reload rebuilds the multipart [text, image_url] content
+ (role-agnostic reconstruction) rather than the flattened text alone."""
+ backend.register_workstream("ws-tool")
+ # assistant(tool_calls) → tool row, mirroring a real read_image turn.
+ import json
- def test_pending_attachments_also_cascade(self, backend):
- backend.register_workstream("ws-cas2")
- pending = _aid()
- consumed = _aid()
- backend.save_attachment(pending, "ws-cas2", "u", "p.txt", "text/plain", 1, "text", b"p")
- backend.save_attachment(consumed, "ws-cas2", "u", "c.txt", "text/plain", 1, "text", b"c")
- msg_id = backend.save_message("ws-cas2", "user", "hi")
- backend.mark_attachments_consumed([consumed], msg_id, "ws-cas2", "u")
+ backend.save_message(
+ "ws-tool",
+ "assistant",
+ None,
+ tool_calls=json.dumps(
+ [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {"name": "read_file", "arguments": "{}"},
+ }
+ ]
+ ),
+ )
+ tool_mid = backend.save_message(
+ "ws-tool", "tool", "Image file: dog.png", "read_file", tool_call_id="c1"
+ )
+ img_id = _hash(PNG_1x1)
+ backend.save_attachment(
+ img_id,
+ "ws-tool",
+ "u",
+ "read_file-image.png",
+ "image/png",
+ len(PNG_1x1),
+ "image",
+ PNG_1x1,
+ "tool",
+ )
+ backend.set_message_attachments("ws-tool", tool_mid, [img_id])
- assert backend.delete_workstream("ws-cas2") is True
- assert backend.get_attachment(pending) is None
- assert backend.get_attachment(consumed) is None
+ msgs = backend.load_messages("ws-tool")
+ tool_msg = next(m for m in msgs if m["role"] == "tool")
+ assert tool_msg["tool_call_id"] == "c1"
+ content = tool_msg["content"]
+ assert isinstance(content, list)
+ assert content[0] == {"type": "text", "text": "Image file: dog.png"}
+ assert content[1]["type"] == "image_url"
+ assert content[1]["image_url"]["url"].startswith("data:image/png;base64,")
+
+ def test_tool_row_without_attachments_stays_string(self, backend):
+ backend.register_workstream("ws-tool2")
+ backend.save_message("ws-tool2", "tool", "plain output", "bash", tool_call_id="c9")
+ tool_msg = backend.load_messages("ws-tool2", repair=False)[0]
+ assert tool_msg["content"] == "plain output"
class TestReconstructMetaSibling:
def test_reconstructed_user_msg_carries_attachments_meta(self, backend):
backend.register_workstream("ws-meta")
- aid = _aid()
+ aid = _hash(b"hi")
backend.save_attachment(aid, "ws-meta", "u", "doc.md", "text/markdown", 2, "text", b"hi")
mid = backend.save_message("ws-meta", "user", "see this")
- backend.mark_attachments_consumed([aid], mid, "ws-meta", "u")
-
- msgs = backend.load_messages("ws-meta")
- assert len(msgs) == 1
- meta = msgs[0].get("_attachments_meta")
+ backend.set_message_attachments("ws-meta", mid, [aid])
+ meta = backend.load_messages("ws-meta")[0].get("_attachments_meta")
assert isinstance(meta, list) and len(meta) == 1
- assert meta[0] == {
- "kind": "text",
- "filename": "doc.md",
- "mime_type": "text/markdown",
- }
+ assert meta[0] == {"kind": "text", "filename": "doc.md", "mime_type": "text/markdown"}
-class TestReservation:
- def test_reserve_excludes_from_pending_listing(self, backend):
- backend.register_workstream("ws-res1")
- aid = _aid()
- backend.save_attachment(aid, "ws-res1", "u", "a.md", "text/plain", 1, "text", b"a")
- assert len(backend.list_pending_attachments("ws-res1", "u")) == 1
- reserved = backend.reserve_attachments([aid], "q-1", "ws-res1", "u")
- assert reserved == [aid]
- # Reserved row must be hidden from the pending list
- assert backend.list_pending_attachments("ws-res1", "u") == []
- # And from the with-content variant used by auto-consume
- assert backend.get_pending_attachments_with_content("ws-res1", "u") == []
-
- def test_reserve_blocks_delete(self, backend):
- backend.register_workstream("ws-res2")
- aid = _aid()
- backend.save_attachment(aid, "ws-res2", "u", "a.md", "text/plain", 1, "text", b"a")
- backend.reserve_attachments([aid], "q-1", "ws-res2", "u")
- # Reserved attachment cannot be deleted — the user must dequeue
- # the queued message first.
- assert backend.delete_attachment(aid, "ws-res2", "u") is False
- assert backend.get_attachment(aid) is not None
-
- def test_reserve_twice_is_idempotent_first_wins(self, backend):
- backend.register_workstream("ws-res3")
- aid = _aid()
- backend.save_attachment(aid, "ws-res3", "u", "a.md", "text/plain", 1, "text", b"a")
- assert backend.reserve_attachments([aid], "q-1", "ws-res3", "u") == [aid]
- # Second reservation for a different queue msg must not steal
- assert backend.reserve_attachments([aid], "q-2", "ws-res3", "u") == []
- row = backend.get_attachment(aid)
- assert row["reserved_for_msg_id"] == "q-1"
-
- def test_unreserve_returns_to_pending(self, backend):
- backend.register_workstream("ws-res4")
- aid = _aid()
- backend.save_attachment(aid, "ws-res4", "u", "a.md", "text/plain", 1, "text", b"a")
- backend.reserve_attachments([aid], "q-1", "ws-res4", "u")
- backend.unreserve_attachments("q-1", "ws-res4", "u")
- # Back to pending — delete and listing work again
- assert len(backend.list_pending_attachments("ws-res4", "u")) == 1
- row = backend.get_attachment(aid)
- assert row["reserved_for_msg_id"] is None
-
- def test_consume_clears_reservation(self, backend):
- backend.register_workstream("ws-res5")
- aid = _aid()
- backend.save_attachment(aid, "ws-res5", "u", "a.md", "text/plain", 1, "text", b"a")
- backend.reserve_attachments([aid], "q-1", "ws-res5", "u")
- mid = backend.save_message("ws-res5", "user", "go")
- backend.mark_attachments_consumed([aid], mid, "ws-res5", "u")
- row = backend.get_attachment(aid)
- # Transition reserved → consumed clears the reservation
- assert row["message_id"] == mid
- assert row["reserved_for_msg_id"] is None
-
- def test_reserve_scoped_to_owner(self, backend):
- backend.register_workstream("ws-res6")
- aid = _aid()
- backend.save_attachment(aid, "ws-res6", "owner", "a.md", "text/plain", 1, "text", b"a")
- # An intruder user_id cannot reserve someone else's attachment
- assert backend.reserve_attachments([aid], "q-x", "ws-res6", "intruder") == []
- row = backend.get_attachment(aid)
- assert row["reserved_for_msg_id"] is None
-
-
-class TestGetAttachmentsRobustness:
- def test_mixed_known_and_unknown_ids(self, backend):
- backend.register_workstream("ws-mix")
- known = _aid()
- unknown = _aid()
- backend.save_attachment(known, "ws-mix", "u", "k.txt", "text/plain", 1, "text", b"k")
- rows = backend.get_attachments([known, unknown, "definitely-not-an-id"])
- assert len(rows) == 1
- assert rows[0]["attachment_id"] == known
-
-
-class TestRewindTruncationCascadesAttachments:
- def test_delete_messages_after_removes_linked_attachments(self, backend):
+class TestRefcountGC:
+ def test_delete_messages_after_decrements_and_prunes(self, backend):
backend.register_workstream("ws-rewind")
- # Two user turns, each with an attachment. A rewind that keeps
- # only the first turn's messages must also drop the second
- # turn's attachment rather than leak the BLOB.
- a1 = _aid()
- a2 = _aid()
- backend.save_attachment(a1, "ws-rewind", "u", "keep.md", "text/plain", 1, "text", b"k")
+ # Two turns, each referencing a distinct blob. A rewind that keeps
+ # only the first turn must prune the second's blob (refcount → 0) and
+ # keep the first's.
+ a1, a2 = _hash(b"keep"), _hash(b"drop")
m1 = backend.save_message("ws-rewind", "user", "turn1")
- backend.mark_attachments_consumed([a1], m1, "ws-rewind", "u")
+ backend.save_attachment(a1, "ws-rewind", "u", "keep.md", "text/plain", 4, "text", b"keep")
+ backend.set_message_attachments("ws-rewind", m1, [a1])
- backend.save_attachment(a2, "ws-rewind", "u", "drop.md", "text/plain", 1, "text", b"d")
m2 = backend.save_message("ws-rewind", "user", "turn2")
- backend.mark_attachments_consumed([a2], m2, "ws-rewind", "u")
+ backend.save_attachment(a2, "ws-rewind", "u", "drop.md", "text/plain", 4, "text", b"drop")
+ backend.set_message_attachments("ws-rewind", m2, [a2])
- # Keep only the first conversation row
- backend.delete_messages_after("ws-rewind", 1)
+ backend.delete_messages_after("ws-rewind", 1) # keep only turn1's row
- # Kept attachment survives
- assert backend.get_attachment(a1) is not None
- # Doomed attachment is gone — no orphan BLOB
- assert backend.get_attachment(a2) is None
+ assert backend.get_attachment(a1) is not None # still referenced
+ assert backend.get_attachment(a2) is None # pruned at refcount 0
- def test_delete_messages_after_preserves_pending(self, backend):
- # Pending (un-consumed) attachments must not be touched by a
- # truncation — they have no message_id and shouldn't be swept
- # up by the cascade.
- backend.register_workstream("ws-rewind2")
- pending = _aid()
- consumed = _aid()
- backend.save_attachment(pending, "ws-rewind2", "u", "p.md", "text/plain", 1, "text", b"p")
- backend.save_attachment(consumed, "ws-rewind2", "u", "c.md", "text/plain", 1, "text", b"c")
- m1 = backend.save_message("ws-rewind2", "user", "turn1")
- backend.mark_attachments_consumed([consumed], m1, "ws-rewind2", "u")
+ def test_deduped_blob_survives_partial_delete(self, backend):
+ """A blob referenced by two messages survives deleting one of them —
+ refcount drops 2 → 1, the blob stays until the last reference goes."""
+ backend.register_workstream("ws-shared")
+ shared = _hash(b"shared-bytes")
+ m1 = backend.save_message("ws-shared", "user", "first")
+ backend.save_attachment(
+ shared, "ws-shared", "u", "s.txt", "text/plain", 12, "text", b"shared-bytes"
+ )
+ backend.set_message_attachments("ws-shared", m1, [shared])
+ m2 = backend.save_message("ws-shared", "user", "second")
+ backend.save_attachment(
+ shared, "ws-shared", "u", "s.txt", "text/plain", 12, "text", b"shared-bytes"
+ )
+ backend.set_message_attachments("ws-shared", m2, [shared])
- backend.delete_messages_after("ws-rewind2", 0) # drop everything
+ assert backend.get_attachment(shared)["refcount"] == 2
+ backend.delete_messages_after("ws-shared", 1) # drop the 2nd turn
+ row = backend.get_attachment(shared)
+ assert row is not None
+ assert row["refcount"] == 1
- # Pending survives (no message_id → no cascade match)
- assert backend.get_attachment(pending) is not None
- # Consumed is dropped with its parent message
- assert backend.get_attachment(consumed) is None
+ def test_delete_workstream_prunes_referenced_blobs(self, backend):
+ backend.register_workstream("ws-cas")
+ aid = _hash(b"a")
+ m = backend.save_message("ws-cas", "user", "hi")
+ backend.save_attachment(aid, "ws-cas", "u", "a.txt", "text/plain", 1, "text", b"a")
+ backend.set_message_attachments("ws-cas", m, [aid])
+
+ assert backend.delete_workstream("ws-cas") is True
+ assert backend.get_attachment(aid) is None
+
+ def test_delete_workstream_keeps_blob_shared_with_other_ws(self, backend):
+ """Content-addressed ids are global: a blob referenced from two
+ workstreams must only be decremented (not blanket-deleted) when one
+ workstream is removed."""
+ backend.register_workstream("ws-one")
+ backend.register_workstream("ws-two")
+ shared = _hash(b"cross-ws")
+ m1 = backend.save_message("ws-one", "user", "a")
+ backend.save_attachment(
+ shared, "ws-one", "u", "s.txt", "text/plain", 8, "text", b"cross-ws"
+ )
+ backend.set_message_attachments("ws-one", m1, [shared])
+ m2 = backend.save_message("ws-two", "user", "b")
+ backend.save_attachment(
+ shared, "ws-two", "u", "s.txt", "text/plain", 8, "text", b"cross-ws"
+ )
+ backend.set_message_attachments("ws-two", m2, [shared])
+ assert backend.get_attachment(shared)["refcount"] == 2
+
+ backend.delete_workstream("ws-one")
+ row = backend.get_attachment(shared)
+ assert row is not None, "blob still referenced by ws-two must survive"
+ assert row["refcount"] == 1
+ backend.delete_workstream("ws-two")
+ assert backend.get_attachment(shared) is None
+
+
+class TestOwnershipGate:
+ def test_referenced_in_ws_true_for_referencing_row(self, backend):
+ backend.register_workstream("ws-own")
+ aid = _hash(b"owned")
+ m = backend.save_message("ws-own", "user", "hi")
+ backend.save_attachment(aid, "ws-own", "u", "o.txt", "text/plain", 5, "text", b"owned")
+ backend.set_message_attachments("ws-own", m, [aid])
+ assert backend.attachment_referenced_in_ws(aid, "ws-own") is True
+
+ def test_referenced_in_ws_false_for_other_ws(self, backend):
+ # The blob is global, but the OTHER workstream has no row referencing
+ # it → the get_content ownership gate denies it there.
+ backend.register_workstream("ws-own2")
+ backend.register_workstream("ws-stranger")
+ aid = _hash(b"owned2")
+ m = backend.save_message("ws-own2", "user", "hi")
+ backend.save_attachment(aid, "ws-own2", "u", "o.txt", "text/plain", 6, "text", b"owned2")
+ backend.set_message_attachments("ws-own2", m, [aid])
+ assert backend.attachment_referenced_in_ws(aid, "ws-stranger") is False
+
+ def test_referenced_in_ws_false_when_unreferenced(self, backend):
+ # A blob written but not yet recorded on any row (shouldn't happen in
+ # the live flow, but the gate must be closed-by-default).
+ backend.register_workstream("ws-own3")
+ aid = _hash(b"orphan")
+ backend.save_attachment(aid, "ws-own3", "u", "o.txt", "text/plain", 6, "text", b"orphan")
+ assert backend.attachment_referenced_in_ws(aid, "ws-own3") is False
@pytest.mark.parametrize("kind", ["image", "text"])
class TestParametrizedKind:
def test_roundtrip_content_bytes(self, backend, kind):
backend.register_workstream(f"ws-p-{kind}")
- aid = _aid()
payload = PNG_1x1 if kind == "image" else b"x" * 42
mime = "image/png" if kind == "image" else "text/plain"
+ aid = _hash(payload)
backend.save_attachment(
aid, f"ws-p-{kind}", "u", f"f.{kind}", mime, len(payload), kind, payload
)
@@ -394,148 +419,3 @@ class TestParametrizedKind:
assert len(rows) == 1
assert rows[0]["content"] == payload
assert rows[0]["kind"] == kind
-
-
-class TestSweepOrphanReservations:
- """Defensive sweep for reservations leaked by process crashes between
- reserve_attachments and consume/unreserve."""
-
- def _backdate(self, backend, attachment_id, *, created_ago=None, reserved_ago=None):
- """Rewrite the row's `created` and/or `reserved_at` columns so the
- sweep sees them as older than they really are.
-
- Works against the same string format the storage layer writes
- (ISO-8601 truncated to seconds).
- """
- from datetime import UTC, datetime, timedelta
-
- import sqlalchemy as sa
-
- from turnstone.core.storage._schema import workstream_attachments
-
- values: dict[str, str] = {}
- if created_ago is not None:
- values["created"] = (datetime.now(UTC) - timedelta(seconds=created_ago)).strftime(
- "%Y-%m-%dT%H:%M:%S"
- )
- if reserved_ago is not None:
- values["reserved_at"] = (datetime.now(UTC) - timedelta(seconds=reserved_ago)).strftime(
- "%Y-%m-%dT%H:%M:%S"
- )
- if not values:
- return
- with backend._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(workstream_attachments.c.attachment_id == attachment_id)
- .values(**values)
- )
- conn.commit()
-
- def test_clears_old_reserved_rows(self, backend):
- backend.register_workstream("ws-sw")
- aid = _aid()
- backend.save_attachment(aid, "ws-sw", "u", "a.txt", "text/plain", 5, "text", b"hello")
- reserved = backend.reserve_attachments([aid], "send-old", "ws-sw", "u")
- assert reserved == [aid]
- # Backdate the reservation timestamp so the sweep considers it stale
- self._backdate(backend, aid, reserved_ago=7200)
-
- n = backend.sweep_orphan_reservations(older_than_seconds=3600)
- assert n == 1
-
- # The row is back in pending — list_pending_attachments will surface it
- pending = backend.list_pending_attachments("ws-sw", "u")
- assert any(p["attachment_id"] == aid for p in pending)
-
- def test_leaves_fresh_reservations_alone(self, backend):
- backend.register_workstream("ws-sw2")
- aid = _aid()
- backend.save_attachment(aid, "ws-sw2", "u", "a.txt", "text/plain", 5, "text", b"hello")
- backend.reserve_attachments([aid], "send-fresh", "ws-sw2", "u")
- # No backdating — reservation was just created
-
- n = backend.sweep_orphan_reservations(older_than_seconds=3600)
- assert n == 0
-
- # Reservation still held
- pending = backend.list_pending_attachments("ws-sw2", "u")
- assert pending == []
-
- def test_old_upload_with_fresh_reservation_is_preserved(self, backend):
- """Regression: an attachment uploaded long ago but reserved just
- now must NOT be swept. ``reserved_at`` (set on reserve) is the
- staleness signal — not ``created`` (upload time)."""
- backend.register_workstream("ws-sw-mix")
- aid = _aid()
- backend.save_attachment(aid, "ws-sw-mix", "u", "a.txt", "text/plain", 5, "text", b"hello")
- # Backdate the upload by a day, but reserve fresh.
- self._backdate(backend, aid, created_ago=86_400)
- reserved = backend.reserve_attachments([aid], "send-fresh", "ws-sw-mix", "u")
- assert reserved == [aid]
-
- n = backend.sweep_orphan_reservations(older_than_seconds=3600)
- assert n == 0
-
- # Reservation still held — pending list is empty
- assert backend.list_pending_attachments("ws-sw-mix", "u") == []
- # And consume against the original send_id still succeeds
- msg_id = backend.save_message("ws-sw-mix", "user", "after fresh reserve")
- backend.mark_attachments_consumed(
- [aid], msg_id, "ws-sw-mix", "u", reserved_for_msg_id="send-fresh"
- )
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["message_id"] == msg_id
-
- def test_consume_clears_reserved_at(self, backend):
- """Once consumed, the row's reservation metadata must be wiped so
- a follow-up sweep can't accidentally match on it."""
- backend.register_workstream("ws-sw-consume")
- aid = _aid()
- backend.save_attachment(
- aid, "ws-sw-consume", "u", "a.txt", "text/plain", 5, "text", b"hello"
- )
- backend.reserve_attachments([aid], "send-c", "ws-sw-consume", "u")
- msg_id = backend.save_message("ws-sw-consume", "user", "consumed")
- backend.mark_attachments_consumed(
- [aid], msg_id, "ws-sw-consume", "u", reserved_for_msg_id="send-c"
- )
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["reserved_at"] is None
- assert row["reserved_for_msg_id"] is None
-
- def test_unreserve_clears_reserved_at(self, backend):
- backend.register_workstream("ws-sw-unres")
- aid = _aid()
- backend.save_attachment(aid, "ws-sw-unres", "u", "a.txt", "text/plain", 5, "text", b"hello")
- backend.reserve_attachments([aid], "send-u", "ws-sw-unres", "u")
- backend.unreserve_attachments("send-u", "ws-sw-unres", "u")
- row = backend.get_attachment(aid)
- assert row is not None
- assert row["reserved_at"] is None
- assert row["reserved_for_msg_id"] is None
-
- def test_skips_consumed_rows(self, backend):
- backend.register_workstream("ws-sw3")
- aid = _aid()
- backend.save_attachment(aid, "ws-sw3", "u", "a.txt", "text/plain", 5, "text", b"hello")
- backend.reserve_attachments([aid], "send-c", "ws-sw3", "u")
- msg_id = backend.save_message("ws-sw3", "user", "consumed turn")
- backend.mark_attachments_consumed(
- [aid], msg_id, "ws-sw3", "u", reserved_for_msg_id="send-c"
- )
- # Even backdating both timestamps shouldn't matter — the sweep
- # excludes consumed rows.
- self._backdate(backend, aid, created_ago=7200, reserved_ago=7200)
-
- n = backend.sweep_orphan_reservations(older_than_seconds=3600)
- assert n == 0
-
- def test_zero_threshold_is_noop(self, backend):
- # Defensive guard against accidental "sweep everything" calls
- n = backend.sweep_orphan_reservations(older_than_seconds=0)
- assert n == 0
- n = backend.sweep_orphan_reservations(older_than_seconds=-5)
- assert n == 0
diff --git a/tests/test_storage_sqlite.py b/tests/test_storage_sqlite.py
index c79e6d3f..9570d626 100644
--- a/tests/test_storage_sqlite.py
+++ b/tests/test_storage_sqlite.py
@@ -217,34 +217,40 @@ class TestLoadMessagesLimit:
assert tail[2]["content"] == "summarized"
def test_limit_bounds_attachment_scan(self, backend):
- """When ``limit=N`` is set, ``load_attachments_for_messages``
- receives only the fetched message ids — the attachment query
- must not fall back to a full-workstream scan. Otherwise the
- tail-N optimization on conversations is partly undone for
- workstreams with many attachments."""
+ """The content-addressed attachment resolution only fetches blobs the
+ *fetched* rows reference — so a tail-N load that doesn't include the
+ attachment-bearing row issues no blob fetch at all, and a full load
+ fetches exactly the referenced ids. This keeps the tail-N conversations
+ LIMIT from being undone by a full-workstream attachment scan."""
from unittest.mock import patch
backend.register_workstream("s1")
+ # Oldest row carries the attachment; then 20 plain rows after it.
+ aid = "a" * 64
+ att_msg_id = backend.save_message("s1", "user", "see attachment")
+ backend.save_attachment(aid, "s1", "u", "x.txt", "text/plain", 1, "text", b"x")
+ backend.set_message_attachments("s1", att_msg_id, [aid])
for i in range(20):
backend.save_message("s1", "user", f"msg-{i:02d}")
- captured: dict[str, list[int] | None] = {}
- orig = backend.load_attachments_for_messages
+ captured: list[list[str]] = []
+ orig = backend.get_attachments
- def _spy(ws_id, *, message_ids=None):
- captured["message_ids"] = list(message_ids) if message_ids is not None else None
- return orig(ws_id, message_ids=message_ids)
+ def _spy(ids):
+ captured.append(sorted(ids))
+ return orig(ids)
- with patch.object(backend, "load_attachments_for_messages", side_effect=_spy):
+ # Tail-N=5 fetches only the 5 newest rows (all plain) — the
+ # attachment row is excluded, so NO blob fetch is issued.
+ with patch.object(backend, "get_attachments", side_effect=_spy):
backend.load_messages("s1", limit=5)
- # Tail-N request passed a bounded list of exactly 5 ids.
- assert captured["message_ids"] is not None
- assert len(captured["message_ids"]) == 5
+ assert captured == []
- with patch.object(backend, "load_attachments_for_messages", side_effect=_spy):
+ # Full load resolves exactly the one referenced id (not a full scan).
+ captured.clear()
+ with patch.object(backend, "get_attachments", side_effect=_spy):
backend.load_messages("s1")
- # Full-load request passes None → backend scans all attachments.
- assert captured["message_ids"] is None
+ assert captured == [[aid]]
class TestSaveMessagesBulk:
diff --git a/turnstone/console/coordinator_adapter.py b/turnstone/console/coordinator_adapter.py
index d5933101..7ebc492d 100644
--- a/turnstone/console/coordinator_adapter.py
+++ b/turnstone/console/coordinator_adapter.py
@@ -23,7 +23,6 @@ from turnstone.core.child_event_bus import ChildEventBus
from turnstone.core.child_source import ClusterChildSource
from turnstone.core.children_registry import ChildrenRegistry
from turnstone.core.log import get_logger
-from turnstone.core.session import AttachmentsNotQueueableError
from turnstone.core.workstream import Workstream, WorkstreamKind, WorkstreamState
if TYPE_CHECKING:
@@ -274,13 +273,13 @@ class CoordinatorAdapter:
Optional ``attachments`` + ``send_id`` carry create-time
attachments onto the first turn dispatched by the lifted
``create`` handler's ``_coord_create_post_install``. The
- send_id token must match the reservation already taken
- against the attachment rows (see
- :func:`turnstone.core.attachments.reserve_and_resolve_attachments`);
- the worker's failure path unreserves so a worker crash
- doesn't leave the rows soft-locked. Both kwargs default to
- ``None`` so the steady-state ``coord_adapter.send`` call
- sites (no attachments) keep working unchanged.
+ attachments were resolved (peeked) from the per-node upload buffer by
+ :func:`turnstone.core.attachments.resolve_staged_attachments`; the
+ committing ``ChatSession.send`` drains them and persists them
+ content-addressed. ``send_id`` is a tracking token only — no
+ reservation to release on a worker crash. Both kwargs default to
+ ``None`` so the steady-state ``coord_adapter.send`` call sites (no
+ attachments) keep working unchanged.
"""
mgr = self._manager
if mgr is None:
@@ -297,29 +296,16 @@ class CoordinatorAdapter:
# mutable kwargs through the call-site frame after return.
_attachments = attachments or None
_send_id = send_id if _attachments else None
- _user_id = ws.user_id
def _run() -> None:
try:
session.send(message, attachments=_attachments, send_id=_send_id)
except Exception:
- # Unreserve any attachments we soft-locked for this
- # send_id so the rows return to pending and don't stay
- # locked forever after a worker crash. Mirrors the
- # interactive create-with-attachments worker pattern.
- if _attachments and _send_id:
- from turnstone.core.memory import (
- unreserve_attachments as _unreserve,
- )
-
- try:
- _unreserve(_send_id, ws_ref.id, _user_id)
- except Exception:
- log.debug(
- "coord_adapter.attachment_unreserve_failed ws=%s",
- ws_ref.id[:8],
- exc_info=True,
- )
+ # Attachments were resolved (peeked) from the per-node upload
+ # buffer, not soft-locked — there is no reservation to release
+ # on a worker crash. Undrained staged bytes expire on the
+ # buffer TTL; the bytes for a turn that DID commit are already
+ # persisted content-addressed.
log.exception("coord_adapter.worker_failed ws=%s", ws_ref.id[:8])
# ``session.send()`` already surfaced the failure to the
# SSE stream (``ui.on_error``), persisted the sanitized
@@ -335,30 +321,14 @@ class CoordinatorAdapter:
# ``AttachmentsNotQueueableError``). The route handler's
# _enqueue catches the rejection and surfaces an
# ``attachments_busy`` status to the caller; the coord
- # adapter's caller has no equivalent return channel, so
- # we mirror the cleanup (release the reservation taken
- # for ``_send_id``) and let session_worker.send return
- # False — the only call site today
- # (``_coord_create_post_install``) hits the spawn branch
- # on a fresh workstream so the catch is defense-in-depth.
+ # adapter's caller has no equivalent return channel, so we let
+ # session_worker.send return False — the only call site today
+ # (``_coord_create_post_install``) hits the spawn branch on a
+ # fresh workstream so the catch is defense-in-depth. Nothing to
+ # release: the staged bytes were peeked, not soft-locked, and a
+ # rejected enqueue never drained them.
att_ids = [a.attachment_id for a in _attachments] if _attachments else None
- try:
- session.queue_message(message, attachment_ids=att_ids, queue_msg_id=_send_id)
- except AttachmentsNotQueueableError:
- if _attachments and _send_id:
- from turnstone.core.memory import (
- unreserve_attachments as _unreserve,
- )
-
- try:
- _unreserve(_send_id, ws_ref.id, _user_id)
- except Exception:
- log.debug(
- "coord_adapter.attachment_unreserve_failed ws=%s",
- ws_ref.id[:8],
- exc_info=True,
- )
- raise
+ session.queue_message(message, attachment_ids=att_ids, queue_msg_id=_send_id)
return session_worker.send(
ws,
diff --git a/turnstone/console/server.py b/turnstone/console/server.py
index 94ff8f74..6c923fe6 100644
--- a/turnstone/console/server.py
+++ b/turnstone/console/server.py
@@ -3383,7 +3383,7 @@ async def _coord_create_post_install(
"""
import uuid as _uuid
- from turnstone.core.attachments import reserve_and_resolve_attachments
+ from turnstone.core.attachments import resolve_staged_attachments
initial_message = (body.get("initial_message") or "").strip()
if not initial_message:
@@ -3392,18 +3392,14 @@ async def _coord_create_post_install(
if coord_adapter is None:
return {}
- # Mirror interactive's reservation pattern: same send_id token
- # scopes the soft-lock and the eventual consume. Coord's
- # ``CoordinatorAdapter.send`` worker passes both through to
- # ``ChatSession.send(..., send_id=...)``; on worker failure the
- # adapter's exception path unreserves so the rows return to
- # pending.
+ # Resolve (peek) the staged uploads for the dispatched first turn; the
+ # committing ``ChatSession.send`` drains them from the per-node buffer and
+ # persists them content-addressed. ``send_id`` is a tracking token only —
+ # no DB reservation to release on worker failure.
send_id = _uuid.uuid4().hex
resolved_atts: list[Any] = []
if attachment_ids:
- resolved_atts, _ord, _drop = reserve_and_resolve_attachments(
- attachment_ids, send_id, ws.id, uid
- )
+ resolved_atts, _ord, _drop = resolve_staged_attachments(attachment_ids, ws.id, uid)
coord_adapter.send(
ws.id,
initial_message,
@@ -12775,14 +12771,10 @@ def create_app(
from turnstone.core.attachments import (
sniff_image_mime as _coord_sniff_image,
)
- from turnstone.core.attachments import (
- upload_lock as _coord_upload_lock,
- )
coord_attachment_helpers = AttachmentUploadHelpers(
sniff_image_mime=_coord_sniff_image,
classify_text_attachment=_coord_classify_text,
- upload_lock=_coord_upload_lock,
)
coord_endpoint_config = SessionEndpointConfig(
permission_gate=_require_admin_coordinator,
diff --git a/turnstone/core/attachments.py b/turnstone/core/attachments.py
index d1585e25..9672b227 100644
--- a/turnstone/core/attachments.py
+++ b/turnstone/core/attachments.py
@@ -1,20 +1,24 @@
"""Attachment data types + upload-classification helpers for user-uploaded files
bound to a workstream turn.
-The image-sniff / text-classify / per-(ws,user) upload-lock helpers
-live here (rather than in ``turnstone/server.py``) so the console
-process can wire them into the lifted attachment endpoints for the
-coordinator surface without depending on the node-side server module.
-The classification policy is intentionally kind-agnostic — the same
-type allowlist applies on both processes.
+The image-sniff / text-classify helpers live here (rather than in
+``turnstone/server.py``) so the console process can wire them into the lifted
+attachment endpoints for the coordinator surface without depending on the
+node-side server module. The classification policy is intentionally
+kind-agnostic — the same type allowlist applies on both processes.
+
+In the content-addressed model an upload is *staged* in the per-node in-memory
+``attachment_buffer`` (keyed by content hash) until the send that references
+it commits, at which point the bytes are written content-addressed +
+reference-counted into ``workstream_attachments``. Staging is thread-safe and
+idempotent on the content hash, so the old per-(ws,user) upload lock + pending
+cap (which only existed to serialize a DB count-check) are gone — the buffer's
+own size/TTL ceilings bound a flood instead.
"""
from __future__ import annotations
-import collections
import os
-import threading
-import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
@@ -25,11 +29,6 @@ if TYPE_CHECKING:
# constants live here so the session / tests share the same definitions.
IMAGE_SIZE_CAP: int = 4 * 1024 * 1024
TEXT_DOC_SIZE_CAP: int = 512 * 1024
-# Cap on simultaneously-pending attachments for a single (ws, user).
-# Once reserved for a queued message the row no longer counts against
-# this budget, so the name reflects the pending-pool limit rather than
-# a per-message limit.
-MAX_PENDING_ATTACHMENTS_PER_USER_WS: int = 10
ALLOWED_IMAGE_MIMES: frozenset[str] = frozenset(
{"image/png", "image/jpeg", "image/gif", "image/webp"}
@@ -60,53 +59,6 @@ class Attachment:
return self.kind == "text"
-# ---------------------------------------------------------------------------
-# Per-(ws, user) upload lock cache
-# ---------------------------------------------------------------------------
-# Soft cap on the upload-lock cache. Locks are evicted opportunistically
-# when an upload completes (see ``upload_lock``); a held lock means an
-# upload is in flight, never evicted.
-_ATTACHMENT_UPLOAD_LOCKS_MAX: int = 1024
-_attachment_upload_locks: collections.OrderedDict[tuple[str, str], threading.Lock] = (
- collections.OrderedDict()
-)
-_attachment_upload_locks_mx: threading.Lock = threading.Lock()
-
-
-def upload_lock(ws_id: str, user_id: str) -> threading.Lock:
- """Return (and track) the per-(ws, user) upload mutex.
-
- Called at the start of every attachment upload to serialize the
- pending-cap check + save sequence per (ws, user) — concurrent
- uploads can't both pass a check that sees ``count == cap-1``.
- Process-local cache; bounded eviction skips held locks.
- """
- key = (ws_id, user_id)
- with _attachment_upload_locks_mx:
- lock = _attachment_upload_locks.get(key)
- if lock is None:
- lock = threading.Lock()
- _attachment_upload_locks[key] = lock
- else:
- # Touch for LRU
- _attachment_upload_locks.move_to_end(key)
- # Opportunistic eviction once we exceed the soft cap. Skip
- # held locks (an upload is in flight under that key).
- if len(_attachment_upload_locks) > _ATTACHMENT_UPLOAD_LOCKS_MAX:
- for stale_key in list(_attachment_upload_locks):
- if len(_attachment_upload_locks) <= _ATTACHMENT_UPLOAD_LOCKS_MAX:
- break
- if stale_key == key:
- continue # never evict the lock we're handing out
- stale = _attachment_upload_locks[stale_key]
- # threading.Lock has no public locked() — use the
- # non-blocking acquire-and-release probe instead.
- if stale.acquire(blocking=False):
- stale.release()
- del _attachment_upload_locks[stale_key]
- return lock
-
-
# ---------------------------------------------------------------------------
# Upload classification
# ---------------------------------------------------------------------------
@@ -197,159 +149,131 @@ def validate_and_save_uploaded_files(
ws_id: str,
user_id: str,
) -> tuple[list[str], JSONResponse | None]:
- """Classify + save a list of ``(filename, claimed_mime, data)`` tuples.
+ """Classify + stage a list of ``(filename, claimed_mime, data)`` tuples.
- Applies the same validation rules as ``upload_attachment`` (magic-byte
- image sniffing, UTF-8 text decode, per-kind size cap, per-(ws,user)
- pending cap) under the shared :func:`upload_lock`.
+ Applies the same validation rules as the upload endpoint (magic-byte image
+ sniffing, UTF-8 text decode, per-kind size cap) and stages each file in the
+ per-node :class:`~turnstone.core.attachment_buffer.AttachmentBuffer`. The
+ returned ids are the content hashes the buffer computed — re-uploading
+ identical bytes is idempotent (same id). The per-(ws,user) pending cap and
+ its lock are gone; the buffer's own size/TTL ceilings bound a flood.
Kind-agnostic: both interactive and coordinator create-with-attachments
paths call into this helper from the lifted ``make_create_handler``
- factory (Stage 2 ``create`` verb lift). The helper does not consult
- any kind-specific config — the storage layer is kind-agnostic by
- design (P1.5).
+ factory.
- Returns ``(attachment_ids, None)`` on success or ``(ids_saved_so_far,
+ Returns ``(attachment_ids, None)`` on success or ``(ids_staged_so_far,
JSONResponse)`` on the first failure so the caller can roll back any
partial state.
"""
from starlette.responses import JSONResponse as _JSONResponse
- from turnstone.core.memory import list_pending_attachments, save_attachment
+ from turnstone.core.attachment_buffer import get_attachment_buffer
saved_ids: list[str] = []
if not files:
return saved_ids, None
- lock = upload_lock(ws_id, user_id)
- with lock:
- pending_count = len(list_pending_attachments(ws_id, user_id))
- for filename, claimed_mime, data in files:
- if not data:
- return saved_ids, _JSONResponse({"error": "Empty file"}, status_code=400)
- sniffed_image = sniff_image_mime(data)
- if sniffed_image is not None:
- if len(data) > IMAGE_SIZE_CAP:
- return saved_ids, _JSONResponse(
- {
- "error": (
- f"Image too large ({len(data):,} bytes); "
- f"cap is {IMAGE_SIZE_CAP:,} bytes."
- ),
- "code": "too_large",
- },
- status_code=413,
- )
- kind = "image"
- mime = sniffed_image
- else:
- if len(data) > TEXT_DOC_SIZE_CAP:
- return saved_ids, _JSONResponse(
- {
- "error": (
- f"Text document too large ({len(data):,} bytes); "
- f"cap is {TEXT_DOC_SIZE_CAP:,} bytes."
- ),
- "code": "too_large",
- },
- status_code=413,
- )
- mime_or_err = classify_text_attachment(filename, claimed_mime, data)
- if mime_or_err[0] is None:
- return saved_ids, _JSONResponse(
- {"error": mime_or_err[1], "code": "unsupported"},
- status_code=400,
- )
- kind = "text"
- mime = mime_or_err[0]
-
- if pending_count + 1 > MAX_PENDING_ATTACHMENTS_PER_USER_WS:
+ buffer = get_attachment_buffer()
+ for filename, claimed_mime, data in files:
+ if not data:
+ return saved_ids, _JSONResponse({"error": "Empty file"}, status_code=400)
+ sniffed_image = sniff_image_mime(data)
+ if sniffed_image is not None:
+ if len(data) > IMAGE_SIZE_CAP:
return saved_ids, _JSONResponse(
{
"error": (
- f"Too many pending attachments "
- f"(max {MAX_PENDING_ATTACHMENTS_PER_USER_WS} pending per workstream)"
+ f"Image too large ({len(data):,} bytes); "
+ f"cap is {IMAGE_SIZE_CAP:,} bytes."
),
- "code": "too_many",
+ "code": "too_large",
},
- status_code=409,
+ status_code=413,
)
- attachment_id = uuid.uuid4().hex
- save_attachment(
- attachment_id,
- ws_id,
- user_id,
- filename,
- mime,
- len(data),
- kind,
- data,
- )
- saved_ids.append(attachment_id)
- pending_count += 1
+ kind = "image"
+ mime = sniffed_image
+ else:
+ if len(data) > TEXT_DOC_SIZE_CAP:
+ return saved_ids, _JSONResponse(
+ {
+ "error": (
+ f"Text document too large ({len(data):,} bytes); "
+ f"cap is {TEXT_DOC_SIZE_CAP:,} bytes."
+ ),
+ "code": "too_large",
+ },
+ status_code=413,
+ )
+ mime_or_err = classify_text_attachment(filename, claimed_mime, data)
+ if mime_or_err[0] is None:
+ return saved_ids, _JSONResponse(
+ {"error": mime_or_err[1], "code": "unsupported"},
+ status_code=400,
+ )
+ kind = "text"
+ mime = mime_or_err[0]
+
+ staged = buffer.stage(
+ ws_id=ws_id,
+ user_id=user_id,
+ filename=filename,
+ mime_type=mime,
+ kind=kind,
+ content=data,
+ )
+ saved_ids.append(staged.attachment_id)
return saved_ids, None
-def reserve_and_resolve_attachments(
+def resolve_staged_attachments(
requested_ids: list[str],
- send_id: str,
ws_id: str,
user_id: str,
) -> tuple[list[Attachment], list[str], list[str]]:
- """Reserve attachment ids for ``send_id`` and resolve to Attachment objects.
+ """Resolve staged uploads for *requested_ids* to Attachment objects.
- Returns ``(resolved, ordered_reserved, dropped)``. ``dropped`` is the
- subset of *requested_ids* that could not be reserved (already consumed,
- lost a race, or cross-scope).
+ Returns ``(resolved, taken, dropped)``. ``taken`` is the subset of
+ *requested_ids* present in the buffer for ``(ws_id, user_id)`` (in request
+ order); ``dropped`` is the rest (buffer-evicted, never staged, or out of
+ scope).
- Kind-agnostic: both interactive and coordinator create-with-attachments
- paths call into this helper from their respective ``post_install``
- callbacks (Stage 2 ``create`` verb lift). The reservation token
- (``send_id``) scopes both the soft-lock and the eventual consume —
- the worker calling ``ChatSession.send(..., send_id=...)`` matches
- the lock and converts pending → consumed; failure paths
- ``unreserve_attachments(send_id, ws_id, user_id)`` to release the
- rows back to pending.
+ This is a *peek*, not a drain: the entries stay in the buffer so a send
+ that resolves them but doesn't commit (e.g. the queue rejects an
+ attachment-bearing turn → ``attachments_busy``, and the client retries) can
+ still find them. The committing path drains them at write time via
+ :meth:`ChatSession._append_user_turn` (``buffer.discard`` per persisted
+ id); anything left over expires on the buffer's TTL.
+
+ Kind-agnostic: both create-with-attachments and ``/send`` paths call this.
+ The old ``send_id`` reservation token is gone — the buffer is the pending
+ store and scoping is ``(ws_id, user_id)`` on the staged entry itself.
"""
- from turnstone.core.memory import get_attachments as _get_attachments
- from turnstone.core.memory import reserve_attachments as _reserve
+ from turnstone.core.attachment_buffer import get_attachment_buffer
if not requested_ids:
return [], [], []
- reserved_ids: list[str] = _reserve(requested_ids, send_id, ws_id, user_id)
- reserved_set = set(reserved_ids)
- ordered_reserved: list[str] = [aid for aid in requested_ids if aid in reserved_set]
- dropped: list[str] = [aid for aid in requested_ids if aid not in reserved_set]
-
+ buffer = get_attachment_buffer()
resolved: list[Attachment] = []
- if ordered_reserved:
- rows = _get_attachments(ordered_reserved)
- rows_by_id = {str(r["attachment_id"]): r for r in rows}
- for aid in ordered_reserved:
- r = rows_by_id.get(aid)
- if not r:
- continue
- if (
- r.get("ws_id") != ws_id
- or r.get("user_id") != user_id
- or r.get("message_id") is not None
- or r.get("reserved_for_msg_id") != send_id
- ):
- continue
- content = r.get("content")
- if not isinstance(content, bytes):
- continue
- resolved.append(
- Attachment(
- attachment_id=str(r["attachment_id"]),
- filename=str(r.get("filename") or ""),
- mime_type=str(r.get("mime_type") or "application/octet-stream"),
- kind=str(r.get("kind") or ""),
- content=content,
- )
+ taken: list[str] = []
+ dropped: list[str] = []
+ for aid in requested_ids:
+ s = buffer.get(aid, ws_id=ws_id, user_id=user_id)
+ if s is None:
+ dropped.append(aid)
+ continue
+ taken.append(aid)
+ resolved.append(
+ Attachment(
+ attachment_id=s.attachment_id,
+ filename=s.filename,
+ mime_type=s.mime_type or "application/octet-stream",
+ kind=s.kind,
+ content=s.content,
)
- return resolved, ordered_reserved, dropped
+ )
+ return resolved, taken, dropped
def unreadable_placeholder(filename: str) -> dict[str, Any]:
diff --git a/turnstone/core/memory.py b/turnstone/core/memory.py
index 117f49ee..efe7833b 100644
--- a/turnstone/core/memory.py
+++ b/turnstone/core/memory.py
@@ -108,8 +108,14 @@ def save_attachment(
size_bytes: int,
kind: str,
content: bytes,
+ origin: str = "upload",
) -> None:
- """Persist an uploaded attachment in pending state."""
+ """Write a content-addressed blob (INSERT-OR-IGNORE) and bump its refcount.
+
+ ``attachment_id`` is the content hash; ``origin`` is ``'upload'`` (user
+ attachment) or ``'tool'`` (e.g. a ``read_file`` image). A blob is only
+ ever written referenced (refcount ≥ 1).
+ """
try:
get_storage().save_attachment(
attachment_id,
@@ -120,18 +126,20 @@ def save_attachment(
size_bytes,
kind,
content,
+ origin,
)
except Exception:
log.warning("Failed to save attachment ws=%s", ws_id, exc_info=True)
-def list_pending_attachments(ws_id: str, user_id: str) -> list[dict[str, Any]]:
- """List un-consumed attachments for ``(ws_id, user_id)``."""
+def set_message_attachments(ws_id: str, message_id: int, attachment_ids: list[str]) -> None:
+ """Record a turn's ordered content-addressed ref-list on its conversations row."""
+ if not attachment_ids or not message_id:
+ return
try:
- return get_storage().list_pending_attachments(ws_id, user_id)
+ get_storage().set_message_attachments(ws_id, message_id, attachment_ids)
except Exception:
- log.warning("Failed to list pending attachments ws=%s", ws_id, exc_info=True)
- return []
+ log.warning("Failed to set message attachments ws=%s", ws_id, exc_info=True)
def get_attachments(attachment_ids: list[str]) -> list[dict[str, Any]]:
@@ -145,22 +153,6 @@ def get_attachments(attachment_ids: list[str]) -> list[dict[str, Any]]:
return []
-def get_pending_attachments_with_content(ws_id: str, user_id: str) -> list[dict[str, Any]]:
- """Single-query fetch of pending attachments + their bytes for the
- auto-consume path on send. Never expose this to user-facing listing
- endpoints — use ``list_pending_attachments`` there instead.
- """
- try:
- return get_storage().get_pending_attachments_with_content(ws_id, user_id)
- except Exception:
- log.warning(
- "Failed to fetch pending attachments with content ws=%s",
- ws_id,
- exc_info=True,
- )
- return []
-
-
def get_attachment(attachment_id: str) -> dict[str, Any] | None:
"""Return a single attachment row (with content) or None."""
try:
@@ -170,97 +162,20 @@ def get_attachment(attachment_id: str) -> dict[str, Any] | None:
return None
-def delete_attachment(attachment_id: str, ws_id: str, user_id: str) -> bool:
- """Delete a pending attachment. Returns True if deleted."""
+def attachment_referenced_in_ws(attachment_id: str, ws_id: str) -> bool:
+ """True iff some conversations row in ``ws_id`` references ``attachment_id``.
+
+ The committed-attachment ownership gate for ``get_content`` (the per-row
+ ws/user scope columns are gone — scope rebases onto referencing-row
+ ownership).
+ """
try:
- return get_storage().delete_attachment(attachment_id, ws_id, user_id)
+ return get_storage().attachment_referenced_in_ws(attachment_id, ws_id)
except Exception:
- log.warning("Failed to delete attachment id=%s", attachment_id, exc_info=True)
+ log.warning("Failed to check attachment reference id=%s", attachment_id, exc_info=True)
return False
-def mark_attachments_consumed(
- attachment_ids: list[str],
- message_id: int,
- ws_id: str,
- user_id: str,
- reserved_for_msg_id: str | None = None,
-) -> None:
- """Link attachments to a saved user message (scoped to ws_id+user_id).
-
- When ``reserved_for_msg_id`` is set, the UPDATE also requires the
- attachment's reservation token to match — prevents a stale send from
- consuming rows reserved for a different one.
- """
- if not attachment_ids:
- return
- try:
- get_storage().mark_attachments_consumed(
- attachment_ids,
- message_id,
- ws_id,
- user_id,
- reserved_for_msg_id=reserved_for_msg_id,
- )
- except Exception:
- log.warning("Failed to mark attachments consumed", exc_info=True)
-
-
-def reserve_attachments(
- attachment_ids: list[str],
- queue_msg_id: str,
- ws_id: str,
- user_id: str,
-) -> list[str]:
- """Soft-lock pending attachments to a queued user message.
-
- Returns the list of ids that were actually reserved for ``queue_msg_id``
- (others silently skipped — e.g. already consumed or reserved).
- """
- if not attachment_ids or not queue_msg_id:
- return []
- try:
- return get_storage().reserve_attachments(attachment_ids, queue_msg_id, ws_id, user_id)
- except Exception:
- log.warning("Failed to reserve attachments", exc_info=True)
- return []
-
-
-def unreserve_attachments(queue_msg_id: str, ws_id: str, user_id: str) -> None:
- """Release the reservation held by ``queue_msg_id`` on this (ws, user)."""
- if not queue_msg_id:
- return
- try:
- get_storage().unreserve_attachments(queue_msg_id, ws_id, user_id)
- except Exception:
- log.warning("Failed to unreserve attachments", exc_info=True)
-
-
-def sweep_orphan_reservations(older_than_seconds: int) -> int:
- """Clear ``reserved_for_msg_id`` on stale attachment rows.
-
- Defensive cleanup for reservations leaked by process crashes between
- ``reserve_attachments`` and ``mark_attachments_consumed`` /
- ``unreserve_attachments``. Returns count of rows swept.
- """
- if older_than_seconds <= 0:
- return 0
- try:
- return get_storage().sweep_orphan_reservations(older_than_seconds)
- except Exception:
- log.warning("Failed to sweep orphan reservations", exc_info=True)
- return 0
-
-
-def load_attachments_for_messages(ws_id: str) -> dict[int, list[dict[str, Any]]]:
- """Return attachments grouped by ``message_id`` for history replay."""
- try:
- return get_storage().load_attachments_for_messages(ws_id)
- except Exception:
- log.warning("Failed to load attachments for ws=%s", ws_id, exc_info=True)
- return {}
-
-
def delete_messages_after(ws_id: str, keep_count: int) -> int:
"""Delete conversation rows beyond the first *keep_count* rows.
diff --git a/turnstone/core/session.py b/turnstone/core/session.py
index dd613e5f..58849c50 100644
--- a/turnstone/core/session.py
+++ b/turnstone/core/session.py
@@ -38,6 +38,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Protocol
import httpx
from turnstone.core import fence
+from turnstone.core.attachment_buffer import get_attachment_buffer
from turnstone.core.attachments import (
IMAGE_SIZE_CAP as _ATTACH_IMAGE_SIZE_CAP,
)
@@ -66,9 +67,9 @@ from turnstone.core.memory import (
list_workstreams_with_history,
load_messages,
load_workstream_config,
- mark_attachments_consumed,
normalize_key,
resolve_workstream,
+ save_attachment,
save_message,
save_messages_bulk,
save_structured_memory,
@@ -77,6 +78,7 @@ from turnstone.core.memory import (
search_history_recent,
search_structured_memories,
search_visible_structured_memories,
+ set_message_attachments,
set_workstream_alias,
update_workstream_title,
)
@@ -3598,14 +3600,17 @@ class ChatSession:
When ``attachments`` is non-empty the in-memory message carries
list content (text + image_url + document parts); the DB
- conversations row stores only the text — attachments link back
- via ``workstream_attachments.message_id``. Returns the saved
- conversations row id (0 on save failure, per the storage
- wrapper's no-raise contract).
+ conversations row stores only the text. At commit each attachment's
+ bytes are written content-addressed + reference-counted into
+ ``workstream_attachments`` (``attachment_id`` = the content hash) and
+ the ordered id-list is recorded on the row's ``attachments`` ref-list
+ column — the sole message->blob link. Returns the saved conversations
+ row id (0 on save failure, per the storage wrapper's no-raise
+ contract).
- ``send_id`` (when provided) is the reservation token; the
- consume step adds it to the WHERE clause so a stale send can't
- steal rows reserved to a different one.
+ ``send_id`` (when provided) is the end-to-end send token; it no longer
+ gates a DB reservation (the upload buffer is the pending store — the
+ bytes were already drained from it before this call).
"""
# New user content invalidates the per-turn memory-search cache
# (composition will see a different recent-context string).
@@ -3683,11 +3688,12 @@ class ChatSession:
]
self.messages.append(user_msg)
self._msg_tokens.append(max(1, int(self._msg_char_count(user_msg) / self._chars_per_token)))
- # DB row stores the raw text only; attachments are joined back in
- # from workstream_attachments on load via message_id. Save →
- # consume are two separate transactions; a crash between them
- # leaves pending rows that the UI's chip rehydration can still
- # surface so the user can clear or resend them.
+ # DB row stores the raw text only; attachment bytes are written
+ # content-addressed into workstream_attachments and the ordered id-list
+ # is recorded on this row's ``attachments`` ref-list column (the sole
+ # message->blob link), joined back in on load. ``send_id`` no longer
+ # gates a reservation — the bytes were drained from the upload buffer
+ # before this call.
#
# The wake's synthesised empty turn carries ``_source`` onto the
# row so reconnecting tabs render the marker instead of an
@@ -3703,14 +3709,90 @@ class ChatSession:
event_id=self._ui_event_id(),
)
if attachments and message_id:
- mark_attachments_consumed(
- [a.attachment_id for a in attachments],
- message_id,
+ self._persist_attachment_refs(message_id, attachments)
+ # Drain the now-committed handles from the per-node upload buffer
+ # (content-addressed: the bytes are persisted + referenced). A
+ # peek-then-commit split (resolve in the route, drain here) lets an
+ # uncommitted send — e.g. one the queue rejected — keep the staged
+ # bytes for a retry; anything not drained expires on the buffer TTL.
+ buffer = get_attachment_buffer()
+ for att in attachments:
+ buffer.discard(att.attachment_id, ws_id=self._ws_id, user_id=self._user_id)
+ return message_id
+
+ def _persist_attachment_refs(
+ self,
+ message_id: int,
+ attachments: list[Attachment] | tuple[Attachment, ...],
+ *,
+ origin: str = "upload",
+ ) -> None:
+ """Write each attachment's bytes content-addressed and record the ref-list.
+
+ ``attachment_id`` is the content hash, so identical bytes dedupe to one
+ blob and each reference bumps its refcount; the ordered id-list is
+ recorded on the conversations row's ``attachments`` column. Used by
+ the user-turn commit (``origin='upload'``) and the tool-image persist
+ (``origin='tool'``).
+ """
+ ref_ids: list[str] = []
+ for att in attachments:
+ save_attachment(
+ att.attachment_id,
self._ws_id,
self._user_id,
- reserved_for_msg_id=send_id,
+ att.filename,
+ att.mime_type,
+ len(att.content),
+ att.kind,
+ att.content,
+ origin,
)
- return message_id
+ ref_ids.append(att.attachment_id)
+ set_message_attachments(self._ws_id, message_id, ref_ids)
+
+ @staticmethod
+ def _image_parts_to_attachments(
+ output: list[dict[str, Any]], tool_name: str
+ ) -> list[Attachment]:
+ """Decode ``image_url`` data-URI parts in tool output into Attachments.
+
+ Tool vision output (e.g. ``read_file`` on an image) carries inline
+ ``data:;base64,<...>`` image parts. Decode them back to bytes,
+ derive the content hash as the ``attachment_id``, and return one
+ ``Attachment`` per image so the caller can persist them
+ content-addressed. Non-image / non-data-URI parts and undecodable
+ payloads are skipped (the inline part still rides the live wire).
+ """
+ atts: list[Attachment] = []
+ for part in output:
+ if not (isinstance(part, dict) and part.get("type") == "image_url"):
+ continue
+ url = (part.get("image_url") or {}).get("url") or ""
+ if not url.startswith("data:") or ";base64," not in url:
+ continue
+ header, _, b64 = url.partition(";base64,")
+ mime = header[len("data:") :] or "image/png"
+ try:
+ # ``binascii.Error`` (raised on a malformed payload) subclasses
+ # ``ValueError``, so a single except covers both.
+ raw = base64.b64decode(b64, validate=True)
+ except ValueError:
+ log.warning(
+ "tool %s emitted an undecodable image data URI; not persisting", tool_name
+ )
+ continue
+ ext = (mimetypes.guess_extension(mime) or ".png").lstrip(".")
+ atts.append(
+ Attachment(
+ attachment_id=hashlib.sha256(raw).hexdigest(),
+ filename=f"{tool_name or 'tool'}-image.{ext}",
+ mime_type=mime,
+ kind="image",
+ content=raw,
+ )
+ )
+ return atts
def _append_system_turn(self, source: str, content: str, **meta: Any) -> None:
"""Append a first-class operator-context system turn and persist it.
@@ -3764,14 +3846,15 @@ class ChatSession:
"""Send user input and handle the response loop (including tool calls).
When ``attachments`` is provided the in-memory user message carries
- multipart list content (text + image_url + document parts) while
- the DB conversations row stores only the text — attachments are
- linked via ``message_id`` in the workstream_attachments table.
+ multipart list content (text + image_url + document parts) while the
+ DB conversations row stores only the text — the attachment bytes are
+ written content-addressed into ``workstream_attachments`` and the
+ ordered id-list is recorded on the row's ``attachments`` ref-list
+ column.
- ``send_id`` is the server-side reservation token for the
- attachments; on consume, the storage layer matches it against
- ``reserved_for_msg_id`` so a stale send can't steal rows
- reserved to a different one.
+ ``send_id`` is an end-to-end tracking token only; it no longer gates a
+ DB reservation (the upload buffer is the pending store, and the bytes
+ in ``attachments`` were already drained/peeked from it by the caller).
"""
self._refresh_model_from_registry()
# Token budget approval gate
@@ -4126,15 +4209,22 @@ class ChatSession:
# already bounded by ``_truncate_output`` above (per-turn
# context budget); no second cap needed.
_tname = _tc_names.get(tc_id, "")
+ tool_image_atts: list[Attachment] = []
if isinstance(output, list):
store_text: str = " ".join(
p.get("text", "")
for p in output
if isinstance(p, dict) and p.get("type") == "text"
)
+ # Persist any image parts content-addressed so vision
+ # tool output (e.g. read_file on an image) survives a
+ # reload — the flattened text alone would drop it. The
+ # in-memory ``output`` keeps its inline image_url for
+ # the live wire; only the ref is persisted.
+ tool_image_atts = self._image_parts_to_attachments(output, _tname)
else:
store_text = output
- save_message(
+ tool_message_id = save_message(
self._ws_id,
"tool",
store_text,
@@ -4143,6 +4233,10 @@ class ChatSession:
event_id=self._ui_event_id(),
is_error=tool_is_error,
)
+ if tool_image_atts and tool_message_id:
+ self._persist_attachment_refs(
+ tool_message_id, tool_image_atts, origin="tool"
+ )
# Accumulate this result's operator context (guard
# findings per-result; queued interjections + metacog
diff --git a/turnstone/core/session_routes.py b/turnstone/core/session_routes.py
index fbd207a6..df391398 100644
--- a/turnstone/core/session_routes.py
+++ b/turnstone/core/session_routes.py
@@ -266,14 +266,14 @@ SavedLoadedLookup = Callable[["Request"], Awaitable[set[str]]]
class AttachmentUploadHelpers:
"""Process-local hooks the lifted attachment factories call into.
- The classification + per-(ws,user) lock are stateful concerns that
- don't belong on the (frozen) :class:`SessionEndpointConfig`
- directly: ``sniff_image_mime`` and ``classify_text_attachment``
- are pure but defined in the kind's owning module;
- ``upload_lock`` returns a process-local cached lock. Bundling
- them on a separate dataclass keeps the cfg declarative and lets
- callers share one helper instance across kinds if the policies
- converge later.
+ The classification helpers are pure but defined in the kind's owning
+ module, so they don't belong on the (frozen)
+ :class:`SessionEndpointConfig` directly. Bundling them on a separate
+ dataclass keeps the cfg declarative and lets callers share one helper
+ instance across kinds if the policies converge later. (The old
+ ``upload_lock`` hook is gone — uploads now stage into the thread-safe,
+ content-addressed per-node buffer, so there's no DB count-check to
+ serialize.)
"""
sniff_image_mime: Callable[[bytes], str | None]
@@ -281,7 +281,6 @@ class AttachmentUploadHelpers:
[str, str, bytes],
tuple[str | None, str | None],
]
- upload_lock: Callable[[str, str], Any]
@dataclass(frozen=True)
@@ -2163,7 +2162,7 @@ def make_create_handler(
uploaded_files: list[tuple[str, str, bytes]] = []
body: dict[str, Any]
if cfg.create_supports_attachments and content_type.startswith("multipart/form-data"):
- # Multipart cap: up to MAX_PENDING × image cap, plus slack
+ # Multipart cap: up to 10 files × the image cap, plus slack
# for JSON meta + multipart framing. Per-file size is
# enforced inside :func:`validate_and_save_uploaded_files`
# against the kind-specific cap.
@@ -3285,12 +3284,12 @@ def make_detail_handler(cfg: SessionEndpointConfig) -> Handler:
def make_send_handler(cfg: SessionEndpointConfig) -> Handler:
"""Lifted body for ``POST {prefix}/{ws_id}/send`` — message dispatch.
- Reserves any attachment ids the request carries, captures a
- ``send_id`` token for end-to-end tracking, then dispatches via
- :func:`turnstone.core.session_worker.send` (atomic
- spawn-or-enqueue under ``ws._lock``). Both queue-reuse and
- spawn paths reserve so the eventual ``mark_attachments_consumed``
- can match on ``reserved_for_msg_id``.
+ Resolves any attachment ids the request carries from the per-node upload
+ buffer (a peek — the bytes stay buffered for a retry if the queue rejects
+ the turn), captures a ``send_id`` tracking token, then dispatches via
+ :func:`turnstone.core.session_worker.send` (atomic spawn-or-enqueue under
+ ``ws._lock``). The committing ``send`` drains the buffer and writes the
+ bytes content-addressed; no reservation is taken or released.
Capability flags on ``cfg`` toggle the kind-specific behaviour:
@@ -3366,29 +3365,17 @@ def make_send_handler(cfg: SessionEndpointConfig) -> Handler:
if ui is None:
return JSONResponse({"error": "session UI not available"}, status_code=409)
- # ----- Attachment reservation (atomic reserve-then-dispatch) -----
+ # ----- Attachment resolution (from the per-node upload buffer) -----
send_id = ""
requested_ids: list[str] = []
ordered_reserved: list[str] = []
reserved_set: set[str] = set()
- reserved_ids: list[str] = []
resolved_atts: list[Any] = []
attach_user_id = ""
if cfg.supports_attachments:
- from turnstone.core.attachments import (
- MAX_PENDING_ATTACHMENTS_PER_USER_WS,
- Attachment,
- )
- from turnstone.core.memory import (
- get_attachments as _get_attachments,
- )
- from turnstone.core.memory import (
- get_pending_attachments_with_content as _get_pending_with_content,
- )
- from turnstone.core.memory import (
- reserve_attachments as _reserve,
- )
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.attachments import resolve_staged_attachments
if cfg.attachment_owner_resolver is None:
# Mis-wired config — the resolver is mandatory when
@@ -3401,79 +3388,32 @@ def make_send_handler(cfg: SessionEndpointConfig) -> Handler:
send_id = uuid.uuid4().hex
raw_ids = body.get("attachment_ids")
- auto_consume_rows: list[dict[str, Any]] = []
if raw_ids is None:
- # Auto-consume: pull the caller's pending (unreserved)
- # rows in creation order — bytes included so we skip
- # a second fetch below.
- auto_consume_rows = _get_pending_with_content(ws_id, attach_user_id)
- requested_ids = [str(r["attachment_id"]) for r in auto_consume_rows]
+ # Auto-consume: every pending (staged) upload for this caller,
+ # in stage order.
+ buffer = get_attachment_buffer()
+ requested_ids = [
+ s.attachment_id for s in buffer.list_for(ws_id=ws_id, user_id=attach_user_id)
+ ]
elif isinstance(raw_ids, list) and raw_ids:
- if len(raw_ids) > MAX_PENDING_ATTACHMENTS_PER_USER_WS:
- return JSONResponse(
- {
- "error": (
- f"Too many attachment_ids "
- f"(max {MAX_PENDING_ATTACHMENTS_PER_USER_WS})"
- ),
- "code": "too_many",
- },
- status_code=400,
- )
requested_ids = [str(x) for x in raw_ids if x]
- reserved_ids = (
- _reserve(requested_ids, send_id, ws_id, attach_user_id) if requested_ids else []
+ # Peek (not drain): the bytes stay buffered so an attachment-bearing
+ # turn the queue rejects can still be retried; the committing
+ # ``send`` drains them at write time. ``resolved`` carries the
+ # bytes the session persists content-addressed.
+ resolved_atts, ordered_reserved, _dropped_resolve = resolve_staged_attachments(
+ requested_ids, ws_id, attach_user_id
)
- reserved_set = set(reserved_ids)
- ordered_reserved = [aid for aid in requested_ids if aid in reserved_set]
-
- if ordered_reserved:
- if auto_consume_rows and all(
- str(r["attachment_id"]) in reserved_set for r in auto_consume_rows
- ):
- rows_by_id = {str(r["attachment_id"]): r for r in auto_consume_rows}
- # reserved_for_msg_id was None at pre-fetch; patch
- # in the token so the scope check below admits the
- # rows we just reserved.
- for r in rows_by_id.values():
- r["reserved_for_msg_id"] = send_id
- else:
- rows = _get_attachments(ordered_reserved)
- rows_by_id = {str(r["attachment_id"]): r for r in rows}
- for aid in ordered_reserved:
- row = rows_by_id.get(aid)
- if not row:
- continue
- # Belt-and-braces scope check on top of the reservation.
- if (
- row.get("ws_id") != ws_id
- or row.get("user_id") != attach_user_id
- or row.get("message_id") is not None
- or row.get("reserved_for_msg_id") != send_id
- ):
- continue
- content = row.get("content")
- if not isinstance(content, bytes):
- continue
- resolved_atts.append(
- Attachment(
- attachment_id=str(row["attachment_id"]),
- filename=str(row.get("filename") or ""),
- mime_type=str(row.get("mime_type") or "application/octet-stream"),
- kind=str(row.get("kind") or ""),
- content=content,
- )
- )
+ reserved_set = set(ordered_reserved)
def _release_reservation_on_fail() -> None:
- """Unreserve if we bail before the dispatcher takes ownership."""
- if reserved_ids:
- from turnstone.core.memory import (
- unreserve_attachments as _unreserve,
- )
+ """No-op: the upload buffer is a peek, not a lock.
- _unreserve(send_id, ws_id, attach_user_id)
+ Retained as the worker-failure hook so the call sites below read
+ the same as the pre-cutover reservation flow; there is nothing to
+ release — undrained staged bytes simply expire on the buffer TTL.
+ """
# If a cancel was just issued, briefly poll for the worker to
# exit before dispatching — avoids spawning into a stale
@@ -3658,7 +3598,6 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
otherwise; they'll just no-op-with-500 when
``attachment_owner_resolver`` is unset.
"""
- import uuid
async def _gate(request: Request) -> JSONResponse | None:
if cfg.permission_gate is not None:
@@ -3677,12 +3616,8 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
return cfg.attachment_owner_resolver(request, ws_id, mgr)
async def upload(request: Request) -> Response:
- from turnstone.core.attachments import (
- IMAGE_SIZE_CAP,
- MAX_PENDING_ATTACHMENTS_PER_USER_WS,
- TEXT_DOC_SIZE_CAP,
- )
- from turnstone.core.memory import list_pending_attachments, save_attachment
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.attachments import IMAGE_SIZE_CAP, TEXT_DOC_SIZE_CAP
from turnstone.core.web_helpers import read_multipart_file_or_400
# Sniffing helpers stay kind-specific because they're tied to
@@ -3692,7 +3627,6 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
return JSONResponse({"error": "attachment_helpers missing"}, status_code=500)
sniff_image = cfg.attachment_helpers.sniff_image_mime
classify_text = cfg.attachment_helpers.classify_text_attachment
- upload_lock = cfg.attachment_helpers.upload_lock
err_gate = await _gate(request)
if err_gate is not None:
@@ -3748,35 +3682,30 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
kind = "text"
mime = mime_or_err[0]
- # Serialize count-check + save per (ws, user) so concurrent
- # uploads can't both pass a check that sees count == cap-1.
- lock = upload_lock(ws_id, user_id)
- with lock:
- if len(list_pending_attachments(ws_id, user_id)) >= MAX_PENDING_ATTACHMENTS_PER_USER_WS:
- return JSONResponse(
- {
- "error": (
- f"Too many pending attachments "
- f"(max {MAX_PENDING_ATTACHMENTS_PER_USER_WS} pending per workstream)"
- ),
- "code": "too_many",
- },
- status_code=409,
- )
- attachment_id = uuid.uuid4().hex
- save_attachment(attachment_id, ws_id, user_id, filename, mime, len(data), kind, data)
+ # Stage in the per-node upload buffer (content-addressed: the id is the
+ # content hash, so re-uploading identical bytes is idempotent). The
+ # bytes are written to storage only at send-commit. No DB write, no
+ # per-user cap — the buffer's size/TTL ceilings bound a flood.
+ staged = get_attachment_buffer().stage(
+ ws_id=ws_id,
+ user_id=user_id,
+ filename=filename,
+ mime_type=mime,
+ kind=kind,
+ content=data,
+ )
return JSONResponse(
{
- "attachment_id": attachment_id,
- "filename": filename,
- "mime_type": mime,
- "size_bytes": len(data),
- "kind": kind,
+ "attachment_id": staged.attachment_id,
+ "filename": staged.filename,
+ "mime_type": staged.mime_type,
+ "size_bytes": staged.size_bytes,
+ "kind": staged.kind,
}
)
async def list_pending(request: Request) -> Response:
- from turnstone.core.memory import list_pending_attachments
+ from turnstone.core.attachment_buffer import get_attachment_buffer
err_gate = await _gate(request)
if err_gate is not None:
@@ -3787,13 +3716,25 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
user_id, err = await _resolve_owner(request, ws_id)
if err:
return err
- rows = list_pending_attachments(ws_id, user_id)
+ # Pending uploads live in the buffer; project to the same wire shape
+ # the DB-backed listing used (no content bytes).
+ rows = [
+ {
+ "attachment_id": s.attachment_id,
+ "filename": s.filename,
+ "mime_type": s.mime_type,
+ "size_bytes": s.size_bytes,
+ "kind": s.kind,
+ }
+ for s in get_attachment_buffer().list_for(ws_id=ws_id, user_id=user_id)
+ ]
return JSONResponse({"attachments": rows})
async def get_content(request: Request) -> Response:
from starlette.responses import Response as _Response
- from turnstone.core.memory import get_attachment
+ from turnstone.core.attachment_buffer import get_attachment_buffer
+ from turnstone.core.memory import attachment_referenced_in_ws, get_attachment
err_gate = await _gate(request)
if err_gate is not None:
@@ -3805,16 +3746,29 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
user_id, err = await _resolve_owner(request, ws_id)
if err:
return err
- row = get_attachment(attachment_id)
- # Scope on user_id too — id-guessing across users in an
- # unowned workstream would otherwise leak blobs. Mask
- # cross-user / cross-ws as 404 to avoid leaking existence.
- if not row or row.get("ws_id") != ws_id or row.get("user_id") != user_id:
- return JSONResponse({"error": "Not found"}, status_code=404)
- body = row.get("content") or b""
- kind = row.get("kind") or ""
- stored_mime = row.get("mime_type") or "application/octet-stream"
- filename = str(row.get("filename") or "attachment")
+
+ # Pending (staged) blobs serve straight from the buffer, scoped to the
+ # uploader. Committed blobs serve from the store, gated by ownership:
+ # the requester (already gated to own ``ws_id``) must have a turn whose
+ # ref-list names the id. Cross-user / cross-ws / unreferenced → 404 so
+ # existence doesn't leak.
+ kind: str
+ stored_mime: str
+ filename: str
+ staged = get_attachment_buffer().get(attachment_id, ws_id=ws_id, user_id=user_id)
+ if staged is not None:
+ body: bytes = staged.content
+ kind = staged.kind
+ stored_mime = staged.mime_type or "application/octet-stream"
+ filename = staged.filename or "attachment"
+ else:
+ row = get_attachment(attachment_id)
+ if not row or not attachment_referenced_in_ws(attachment_id, ws_id):
+ return JSONResponse({"error": "Not found"}, status_code=404)
+ body = row.get("content") or b""
+ kind = row.get("kind") or ""
+ stored_mime = row.get("mime_type") or "application/octet-stream"
+ filename = str(row.get("filename") or "attachment")
# Force text/plain for text kinds — avoids same-origin HTML/SVG
# rendering if a user uploaded an HTML-ish text file. Images
# keep their sniffed MIME (allowlist is strict: png/jpeg/gif/webp).
@@ -3829,7 +3783,7 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
return _Response(body, media_type=response_mime, headers=headers)
async def delete_(request: Request) -> Response:
- from turnstone.core.memory import delete_attachment as _delete
+ from turnstone.core.attachment_buffer import get_attachment_buffer
err_gate = await _gate(request)
if err_gate is not None:
@@ -3841,7 +3795,9 @@ def make_attachment_handlers(cfg: SessionEndpointConfig) -> AttachmentHandlers:
user_id, err = await _resolve_owner(request, ws_id)
if err:
return err
- deleted = _delete(attachment_id, ws_id, user_id)
+ # Only pending (staged) uploads are deletable — a committed blob is
+ # owned by the turn that references it and is GC'd by refcount.
+ deleted = get_attachment_buffer().discard(attachment_id, ws_id=ws_id, user_id=user_id)
if not deleted:
return JSONResponse({"error": "Not found"}, status_code=404)
return JSONResponse({"status": "deleted"})
diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py
index d56f5d6a..43efef2c 100644
--- a/turnstone/core/storage/_postgresql.py
+++ b/turnstone/core/storage/_postgresql.py
@@ -3,14 +3,16 @@
from __future__ import annotations
import contextlib
+import json
import os
import threading
import time
+from collections import Counter
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
- from collections.abc import Iterable, Iterator
+ from collections.abc import Iterable, Iterator, Sequence
from turnstone.core.storage._notify import Notify, NotifyStream
@@ -106,12 +108,18 @@ from turnstone.core.storage._utils import (
from turnstone.core.storage._utils import (
VERDICT_MUTABLE as _VERDICT_MUTABLE,
)
+from turnstone.core.storage._utils import (
+ build_attachments_by_msg as _build_attachments_by_msg,
+)
from turnstone.core.storage._utils import (
escape_like as _escape_like,
)
from turnstone.core.storage._utils import (
normalize_search_terms as _normalize_search_terms,
)
+from turnstone.core.storage._utils import (
+ parse_attachment_refs as _parse_attachment_refs,
+)
from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text
from turnstone.core.storage._utils import (
reconstruct_messages as _reconstruct_messages,
@@ -364,21 +372,27 @@ class PostgreSQLBackend:
def load_messages(
self, ws_id: str, *, limit: int | None = None, repair: bool = True
) -> list[dict[str, Any]]:
+ # The trailing ``attachments`` column carries the per-row
+ # content-addressed ref-list; it is split off below to resolve blobs
+ # and is NOT part of the positional tuple ``reconstruct_messages``
+ # unpacks (id..is_error).
+ _cols = (
+ conversations.c.id,
+ conversations.c.role,
+ conversations.c.content,
+ conversations.c.tool_name,
+ conversations.c.tool_call_id,
+ conversations.c.provider_data,
+ conversations.c.tool_calls,
+ conversations.c._source,
+ conversations.c.event_id,
+ conversations.c.is_error,
+ conversations.c.attachments,
+ )
with self._conn() as conn:
if limit is not None and limit > 0:
rows = conn.execute(
- sa.select(
- conversations.c.id,
- conversations.c.role,
- conversations.c.content,
- conversations.c.tool_name,
- conversations.c.tool_call_id,
- conversations.c.provider_data,
- conversations.c.tool_calls,
- conversations.c._source,
- conversations.c.event_id,
- conversations.c.is_error,
- )
+ sa.select(*_cols)
.where(conversations.c.ws_id == ws_id)
.order_by(conversations.c.id.desc())
.limit(limit)
@@ -386,30 +400,36 @@ class PostgreSQLBackend:
rows = list(reversed(rows))
else:
rows = conn.execute(
- sa.select(
- conversations.c.id,
- conversations.c.role,
- conversations.c.content,
- conversations.c.tool_name,
- conversations.c.tool_call_id,
- conversations.c.provider_data,
- conversations.c.tool_calls,
- conversations.c._source,
- conversations.c.event_id,
- conversations.c.is_error,
- )
+ sa.select(*_cols)
.where(conversations.c.ws_id == ws_id)
.order_by(conversations.c.id)
).fetchall()
- # Bound the attachment scan to the fetched message ids when
- # tail-N was requested — otherwise the attachments query
- # still scans every row for the workstream and partially
- # defeats the conversations-table LIMIT.
- message_ids: list[int] | None = None
- if limit is not None and limit > 0:
- message_ids = [r[0] for r in rows]
- attachments = self.load_attachments_for_messages(ws_id, message_ids=message_ids)
- return _reconstruct_messages(list(rows), ws_id, attachments or None, repair=repair)
+ attachments = self._resolve_row_attachments(rows)
+ # Strip the trailing ref-list column so the tuple shape stays exactly
+ # what ``reconstruct_messages`` expects (id..is_error).
+ msg_rows = [tuple(r)[:10] for r in rows]
+ return _reconstruct_messages(msg_rows, ws_id, attachments or None, repair=repair)
+
+ def _resolve_row_attachments(self, rows: Sequence[Any]) -> dict[int, list[dict[str, Any]]]:
+ """Build the ``reconstruct_messages`` attachment map from row ref-lists.
+
+ Each row's trailing ``attachments`` column (last element) is the
+ content-addressed ref-list; collect every referenced id, bulk-fetch
+ the blobs in one query, and group them back per row id in ref-list
+ order. No referenced ids → no query.
+ """
+ attachment_refs: dict[int, list[str]] = {}
+ all_ids: set[str] = set()
+ for r in rows:
+ ids = _parse_attachment_refs(r[10])
+ if ids:
+ attachment_refs[r[0]] = ids
+ all_ids.update(ids)
+ if not all_ids:
+ return {}
+ blobs = self.get_attachments(list(all_ids))
+ rows_by_id = {str(b["attachment_id"]): b for b in blobs}
+ return _build_attachments_by_msg(attachment_refs, rows_by_id)
def get_max_event_id(self, ws_id: str) -> int | None:
with self._conn() as conn:
@@ -432,16 +452,23 @@ class PostgreSQLBackend:
if cutoff_row is None:
return 0
cutoff_id = cutoff_row[0]
- # Cascade-delete attachments linked to doomed messages so
- # rewind/retry flows don't leak orphan BLOBs.
- conn.execute(
- sa.delete(workstream_attachments).where(
+ # Refcount GC: read the doomed rows' content-addressed ref-lists,
+ # decrement each blob's refcount once per reference, and prune
+ # blobs that hit 0 — so a deduped blob still referenced by a kept
+ # turn survives. Replaces the old message_id-cascade delete.
+ doomed = conn.execute(
+ sa.select(conversations.c.attachments).where(
sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.message_id >= cutoff_id,
+ conversations.c.ws_id == ws_id,
+ conversations.c.id >= cutoff_id,
+ conversations.c.attachments.is_not(None),
)
)
- )
+ ).fetchall()
+ doomed_ids: list[str] = []
+ for (refs,) in doomed:
+ doomed_ids.extend(_parse_attachment_refs(refs))
+ self._release_attachment_refs(conn, doomed_ids)
result = conn.execute(
sa.delete(conversations).where(
sa.and_(
@@ -827,9 +854,22 @@ class PostgreSQLBackend:
def delete_workstream(self, ws_id: str) -> bool:
with self._conn() as conn:
- conn.execute(
- sa.delete(workstream_attachments).where(workstream_attachments.c.ws_id == ws_id)
- )
+ # Refcount GC over every referenced blob (content-addressed ids are
+ # global, so a deduped blob may be shared with another workstream —
+ # decrement, don't blanket-delete by ws_id). Blobs that hit 0 are
+ # pruned; any still referenced elsewhere survive.
+ referenced = conn.execute(
+ sa.select(conversations.c.attachments).where(
+ sa.and_(
+ conversations.c.ws_id == ws_id,
+ conversations.c.attachments.is_not(None),
+ )
+ )
+ ).fetchall()
+ ref_ids: list[str] = []
+ for (refs,) in referenced:
+ ref_ids.extend(_parse_attachment_refs(refs))
+ self._release_attachment_refs(conn, ref_ids)
conn.execute(sa.delete(conversations).where(conversations.c.ws_id == ws_id))
conn.execute(sa.delete(workstream_config).where(workstream_config.c.ws_id == ws_id))
conn.execute(
@@ -845,7 +885,7 @@ class PostgreSQLBackend:
conn.commit()
return result.rowcount > 0
- # -- Workstream attachments ------------------------------------------------
+ # -- Workstream attachments (content-addressed, refcounted) ----------------
def save_attachment(
self,
@@ -857,48 +897,63 @@ class PostgreSQLBackend:
size_bytes: int,
kind: str,
content: bytes,
+ origin: str = "upload",
) -> None:
+ """Write a content-addressed blob (INSERT-OR-IGNORE) and ``refcount += 1``.
+
+ Symmetric with the SQLite backend (see its docstring): the content
+ hash is the PK, the first reference writes at ``refcount = 1`` and
+ subsequent references only bump the count, so a stored blob is always
+ referenced and dedupes across messages / workstreams.
+ """
+ from sqlalchemy.dialects.postgresql import insert as pg_insert
+
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
with self._conn() as conn:
+ stmt = pg_insert(workstream_attachments).values(
+ attachment_id=attachment_id,
+ ws_id=ws_id,
+ user_id=user_id,
+ filename=filename,
+ mime_type=mime_type,
+ size_bytes=size_bytes,
+ kind=kind,
+ content=content,
+ created=now,
+ refcount=0,
+ origin=origin,
+ )
+ conn.execute(stmt.on_conflict_do_nothing(index_elements=["attachment_id"]))
conn.execute(
- sa.insert(workstream_attachments),
- {
- "attachment_id": attachment_id,
- "ws_id": ws_id,
- "user_id": user_id,
- "filename": filename,
- "mime_type": mime_type,
- "size_bytes": size_bytes,
- "kind": kind,
- "content": content,
- "message_id": None,
- "created": now,
- },
+ sa.update(workstream_attachments)
+ .where(workstream_attachments.c.attachment_id == attachment_id)
+ .values(refcount=workstream_attachments.c.refcount + 1)
)
conn.commit()
- def list_pending_attachments(self, ws_id: str, user_id: str) -> list[dict[str, Any]]:
+ def set_message_attachments(
+ self, ws_id: str, message_id: int, attachment_ids: list[str]
+ ) -> None:
+ """Record a turn's ordered content-addressed ref-list on its row.
+
+ Symmetric with the SQLite backend: writes the JSON id-list onto
+ ``conversations.attachments`` for the ``(ws_id, message_id)`` row.
+ Empty input is a no-op.
+ """
+ if not attachment_ids or not message_id:
+ return
with self._conn() as conn:
- rows = conn.execute(
- sa.select(
- workstream_attachments.c.attachment_id,
- workstream_attachments.c.filename,
- workstream_attachments.c.mime_type,
- workstream_attachments.c.size_bytes,
- workstream_attachments.c.kind,
- workstream_attachments.c.created,
- )
+ conn.execute(
+ sa.update(conversations)
.where(
sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
+ conversations.c.id == message_id,
+ conversations.c.ws_id == ws_id,
)
)
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- return [dict(r._mapping) for r in rows]
+ .values(attachments=json.dumps(list(attachment_ids)))
+ )
+ conn.commit()
def get_attachments(self, attachment_ids: list[str]) -> list[dict[str, Any]]:
if not attachment_ids:
@@ -911,24 +966,6 @@ class PostgreSQLBackend:
).fetchall()
return [dict(r._mapping) for r in rows]
- def get_pending_attachments_with_content(
- self, ws_id: str, user_id: str
- ) -> list[dict[str, Any]]:
- with self._conn() as conn:
- rows = conn.execute(
- sa.select(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
- )
- )
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- return [dict(r._mapping) for r in rows]
-
def get_attachment(self, attachment_id: str) -> dict[str, Any] | None:
with self._conn() as conn:
row = conn.execute(
@@ -938,157 +975,53 @@ class PostgreSQLBackend:
).fetchone()
return dict(row._mapping) if row else None
- def delete_attachment(self, attachment_id: str, ws_id: str, user_id: str) -> bool:
+ def attachment_referenced_in_ws(self, attachment_id: str, ws_id: str) -> bool:
+ """True iff some conversations row in ``ws_id`` references ``attachment_id``.
+
+ The committed-attachment ownership gate (see the SQLite sibling for
+ the full rationale): a quoted-id JSON-array substring match on the
+ ``attachments`` column; 64-char sha256 ids cannot collide.
+ """
+ needle = f'%"{attachment_id}"%'
with self._conn() as conn:
- result = conn.execute(
- sa.delete(workstream_attachments).where(
+ row = conn.execute(
+ sa.select(conversations.c.id)
+ .where(
sa.and_(
- workstream_attachments.c.attachment_id == attachment_id,
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
+ conversations.c.ws_id == ws_id,
+ conversations.c.attachments.is_not(None),
+ conversations.c.attachments.like(needle),
)
)
- )
- conn.commit()
- return result.rowcount > 0
+ .limit(1)
+ ).fetchone()
+ return row is not None
- def mark_attachments_consumed(
- self,
- attachment_ids: list[str],
- message_id: int,
- ws_id: str,
- user_id: str,
- reserved_for_msg_id: str | None = None,
- ) -> None:
+ @staticmethod
+ def _release_attachment_refs(conn: Any, attachment_ids: list[str]) -> None:
+ """Decrement refcount once per id and prune blobs that reach 0.
+
+ Symmetric with the SQLite backend: counts duplicate ids in the input
+ so a batch spanning several turns that each reference the same deduped
+ blob decrements by the right amount. Caller holds the transaction.
+ """
if not attachment_ids:
return
- predicate = sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
+ counts = Counter(attachment_ids)
+ for aid, n in counts.items():
+ conn.execute(
+ sa.update(workstream_attachments)
+ .where(workstream_attachments.c.attachment_id == aid)
+ .values(refcount=workstream_attachments.c.refcount - n)
+ )
+ conn.execute(
+ sa.delete(workstream_attachments).where(
+ sa.and_(
+ workstream_attachments.c.attachment_id.in_(list(counts)),
+ workstream_attachments.c.refcount <= 0,
+ )
+ )
)
- if reserved_for_msg_id is not None:
- predicate = sa.and_(
- predicate,
- workstream_attachments.c.reserved_for_msg_id == reserved_for_msg_id,
- )
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(predicate)
- .values(
- message_id=message_id,
- reserved_for_msg_id=None,
- reserved_at=None,
- )
- )
- conn.commit()
-
- def reserve_attachments(
- self,
- attachment_ids: list[str],
- queue_msg_id: str,
- ws_id: str,
- user_id: str,
- ) -> list[str]:
- if not attachment_ids or not queue_msg_id:
- return []
- now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
- )
- )
- .values(reserved_for_msg_id=queue_msg_id, reserved_at=now)
- )
- rows = conn.execute(
- sa.select(workstream_attachments.c.attachment_id).where(
- sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.reserved_for_msg_id == queue_msg_id,
- )
- )
- ).fetchall()
- conn.commit()
- return [r[0] for r in rows]
-
- def unreserve_attachments(self, queue_msg_id: str, ws_id: str, user_id: str) -> None:
- if not queue_msg_id:
- return
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.reserved_for_msg_id == queue_msg_id,
- )
- )
- .values(reserved_for_msg_id=None, reserved_at=None)
- )
- conn.commit()
-
- def sweep_orphan_reservations(self, older_than_seconds: int) -> int:
- if older_than_seconds <= 0:
- return 0
- cutoff = (datetime.now(UTC) - timedelta(seconds=older_than_seconds)).strftime(
- "%Y-%m-%dT%H:%M:%S"
- )
- with self._conn() as conn:
- result = conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.reserved_for_msg_id.is_not(None),
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_at.is_not(None),
- workstream_attachments.c.reserved_at < cutoff,
- )
- )
- .values(reserved_for_msg_id=None, reserved_at=None)
- )
- conn.commit()
- return int(result.rowcount or 0)
-
- def load_attachments_for_messages(
- self,
- ws_id: str,
- *,
- message_ids: list[int] | None = None,
- ) -> dict[int, list[dict[str, Any]]]:
- with self._conn() as conn:
- where_clauses = [
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.message_id.is_not(None),
- ]
- if message_ids is not None:
- if not message_ids:
- return {}
- where_clauses.append(workstream_attachments.c.message_id.in_(message_ids))
- rows = conn.execute(
- sa.select(workstream_attachments)
- .where(sa.and_(*where_clauses))
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- grouped: dict[int, list[dict[str, Any]]] = {}
- for r in rows:
- row = dict(r._mapping)
- mid = row["message_id"]
- grouped.setdefault(mid, []).append(row)
- return grouped
def list_workstreams(
self,
diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py
index 79603663..b86e2b79 100644
--- a/turnstone/core/storage/_protocol.py
+++ b/turnstone/core/storage/_protocol.py
@@ -228,7 +228,15 @@ class StorageBackend(Protocol):
"""
...
- # -- Workstream attachments -----------------------------------------------
+ # -- Workstream attachments (content-addressed, refcounted) ---------------
+ #
+ # Pending (uploaded-but-unsent) bytes live in the per-node in-memory
+ # ``attachment_buffer``, NOT in storage — the persisted pending/reserved/
+ # consumed lifecycle (and its orphan-sweep) was retired by the
+ # content-addressing cutover. Storage holds only committed blobs: written
+ # content-addressed at send-commit (or when a tool produces an image),
+ # deduped by content hash, and reference-counted via the ordered
+ # ``conversations.attachments`` ref-list.
def save_attachment(
self,
@@ -240,15 +248,29 @@ class StorageBackend(Protocol):
size_bytes: int,
kind: str,
content: bytes,
+ origin: str = "upload",
) -> None:
- """Persist an uploaded attachment in pending (unconsumed) state."""
+ """Write a content-addressed blob (INSERT-OR-IGNORE) and ``refcount += 1``.
+
+ ``attachment_id`` is the content hash (sha256 hex); the caller computes
+ it. The first reference inserts the row at ``refcount = 1``; every
+ later reference (a re-upload of identical bytes, or a second message
+ referencing the same blob) only bumps the count. A stored blob is thus
+ always referenced (born at ≥ 1) and identical bytes dedupe to one row
+ across messages and workstreams. ``origin`` is ``'upload'`` (user
+ attachment) or ``'tool'`` (e.g. a ``read_file`` image).
+ """
...
- def list_pending_attachments(self, ws_id: str, user_id: str) -> list[dict[str, Any]]:
- """Return un-consumed attachments for ``(ws_id, user_id)``.
+ def set_message_attachments(
+ self, ws_id: str, message_id: int, attachment_ids: list[str]
+ ) -> None:
+ """Record a turn's ordered content-addressed ref-list on its row.
- Each dict contains: ``attachment_id``, ``filename``, ``mime_type``,
- ``size_bytes``, ``kind``, ``created``. Content bytes are NOT returned.
+ Writes the JSON id-list onto ``conversations.attachments`` for the
+ ``(ws_id, message_id)`` conversations row — the sole message->blob
+ link. Empty input is a no-op (the column stays NULL). Scoped to
+ ``ws_id`` as defense-in-depth against a cross-ws message id.
"""
...
@@ -259,108 +281,18 @@ class StorageBackend(Protocol):
"""
...
- def get_pending_attachments_with_content(
- self, ws_id: str, user_id: str
- ) -> list[dict[str, Any]]:
- """Fetch all pending attachments for ``(ws_id, user_id)`` in a single
- query, including ``content`` bytes.
-
- Used by the auto-consume path on send — saves the two-roundtrip
- list-then-get dance. Excluded by design from the user-facing
- listing API (which must never expose bytes).
- """
- ...
-
def get_attachment(self, attachment_id: str) -> dict[str, Any] | None:
"""Return a single attachment row (with content bytes) or None."""
...
- def delete_attachment(self, attachment_id: str, ws_id: str, user_id: str) -> bool:
- """Delete a pending attachment.
+ def attachment_referenced_in_ws(self, attachment_id: str, ws_id: str) -> bool:
+ """True iff some conversations row in ``ws_id`` references ``attachment_id``.
- Only succeeds when the row matches ``ws_id``, ``user_id``, AND
- ``message_id IS NULL`` (i.e. not yet consumed). Returns True if
- a row was deleted.
- """
- ...
-
- def mark_attachments_consumed(
- self,
- attachment_ids: list[str],
- message_id: int,
- ws_id: str,
- user_id: str,
- reserved_for_msg_id: str | None = None,
- ) -> None:
- """Link a set of attachments to a freshly-saved user message.
-
- The UPDATE is scoped to ``(ws_id, user_id)`` and
- ``message_id IS NULL`` as defense-in-depth: even if a caller
- passes attachment ids that don't belong to them, nothing will be
- consumed. When ``reserved_for_msg_id`` is set, also requires
- the reservation to match — prevents a stale send from consuming
- rows reserved to a different one. Clears ``reserved_for_msg_id``
- on transition.
- """
- ...
-
- def reserve_attachments(
- self,
- attachment_ids: list[str],
- queue_msg_id: str,
- ws_id: str,
- user_id: str,
- ) -> list[str]:
- """Soft-lock pending attachments to a queued user message.
-
- Only rows where ``(ws_id, user_id)`` match and both
- ``message_id`` and ``reserved_for_msg_id`` are NULL are updated.
- Returns the list of ids that were actually reserved (others
- silently skipped — caller should not assume completeness).
- """
- ...
-
- def unreserve_attachments(self, queue_msg_id: str, ws_id: str, user_id: str) -> None:
- """Release any reservation for ``queue_msg_id``.
-
- Used when a queued message is dequeued (cancelled) before
- dispatch — the attachments return to ``pending``.
- """
- ...
-
- def sweep_orphan_reservations(self, older_than_seconds: int) -> int:
- """Clear ``reserved_for_msg_id`` on stale reservations.
-
- Targets rows with ``reserved_for_msg_id IS NOT NULL`` AND
- ``message_id IS NULL`` AND ``reserved_at`` older than the cutoff.
- Self-heals reservations leaked by process crashes between
- ``reserve_attachments`` and ``mark_attachments_consumed`` /
- ``unreserve_attachments``.
-
- Uses ``reserved_at`` (set on reserve, cleared on consume /
- unreserve) rather than ``created`` (upload time) so an attachment
- that sat pending for hours before being reserved is not
- mistakenly unreserved mid-send. Returns the row count swept.
- """
- ...
-
- def load_attachments_for_messages(
- self,
- ws_id: str,
- *,
- message_ids: list[int] | None = None,
- ) -> dict[int, list[dict[str, Any]]]:
- """Return attachments grouped by ``message_id`` for history replay.
-
- Each attachment dict includes ``attachment_id``, ``filename``,
- ``mime_type``, ``size_bytes``, ``kind``, and ``content`` (bytes).
- Pending (un-consumed) rows are excluded.
-
- ``message_ids`` narrows the scan to attachments tied to the
- given message rows — used by the tail-N path in
- :func:`load_messages` so the attachment read doesn't defeat
- the conversations-table LIMIT. Default ``None`` returns every
- attachment for the workstream.
+ The committed-attachment ownership gate: the per-row ``ws_id`` /
+ ``user_id`` scope columns are gone, so ``get_content`` for a committed
+ blob is authorised by proving the requester (already gated to own
+ ``ws_id``) has a turn in that workstream whose ``attachments`` ref-list
+ names the id.
"""
...
diff --git a/turnstone/core/storage/_schema.py b/turnstone/core/storage/_schema.py
index 0bc0d9f2..e13e3108 100644
--- a/turnstone/core/storage/_schema.py
+++ b/turnstone/core/storage/_schema.py
@@ -531,13 +531,24 @@ sa.Index(
)
# ---------------------------------------------------------------------------
-# Workstream attachments — user-uploaded images and text documents bound to
-# a specific user turn (one-shot, consumed when linked to a conversations row).
+# Workstream attachments — content-addressed, refcounted blob store.
+#
+# In the content-addressed model the primary key IS the content hash
+# (sha256 hex): identical bytes dedupe to one row regardless of how many
+# messages reference them. A blob is written only at send-commit (or when a
+# tool produces an image), so every stored row is born referenced
+# (``refcount >= 1``); GC decrements ``refcount`` as referencing messages are
+# deleted and prunes the row at 0. Pending (uploaded-but-unsent) bytes live
+# in the per-node in-memory buffer (``attachment_buffer``), NOT here — the
+# persisted pending/reserved/consumed lifecycle (message_id / reserved_* and
+# its orphan-sweep) was retired by the content-addressing cutover. The
+# message->blob link is the ordered ``conversations.attachments`` ref-list.
# ---------------------------------------------------------------------------
workstream_attachments = sa.Table(
"workstream_attachments",
metadata,
+ # PK is the content hash (sha256 hex) — content-addressed dedup.
sa.Column("attachment_id", sa.Text, primary_key=True),
sa.Column("ws_id", sa.Text, nullable=False),
sa.Column("user_id", sa.Text, nullable=False),
@@ -546,48 +557,15 @@ workstream_attachments = sa.Table(
sa.Column("size_bytes", sa.Integer, nullable=False),
sa.Column("kind", sa.Text, nullable=False), # 'image' | 'text'
sa.Column("content", sa.LargeBinary, nullable=False),
- sa.Column("message_id", sa.Integer, nullable=True), # conversations.id once consumed
- # Soft lock tying an attachment to a queued user message. Lifecycle:
- # pending : message_id IS NULL AND reserved_for_msg_id IS NULL
- # reserved : message_id IS NULL AND reserved_for_msg_id =
- # consumed : message_id IS NOT NULL (reservation cleared on transition)
- sa.Column("reserved_for_msg_id", sa.Text, nullable=True),
- # When the row last transitioned into reserved state. Cleared on
- # consume / unreserve. Set independently of `created` (upload time)
- # so the orphan-reservation sweep can target only reservations that
- # have actually been held longer than the threshold.
- sa.Column("reserved_at", sa.Text, nullable=True),
sa.Column("created", sa.Text, nullable=False),
- # Content-addressed blob store (canonical-trajectory cut): a deduped blob's
- # live-reference count (pruned at 0) and its origin ('upload' | 'tool'). The
- # cutover retires the message_id / reserved_* upload-lifecycle columns above in
- # favour of refcount + the conversations.attachments ref-list.
+ # A deduped blob's live-reference count (pruned at 0) and its origin
+ # ('upload' | 'tool'). The sole message->blob link is the ordered
+ # ``conversations.attachments`` ref-list, not a column here.
sa.Column("refcount", sa.Integer, nullable=False, server_default=sa.text("0")),
sa.Column("origin", sa.Text, nullable=False, server_default=sa.text("'upload'")),
)
sa.Index("idx_ws_attachments_ws_id", workstream_attachments.c.ws_id)
-sa.Index(
- "idx_ws_attachments_pending",
- workstream_attachments.c.ws_id,
- workstream_attachments.c.user_id,
- workstream_attachments.c.message_id,
-)
-sa.Index("idx_ws_attachments_message", workstream_attachments.c.message_id)
-sa.Index(
- "idx_ws_attachments_reserved",
- workstream_attachments.c.ws_id,
- workstream_attachments.c.user_id,
- workstream_attachments.c.reserved_for_msg_id,
-)
-# Partial index — only reserved rows participate, so the sweep scan
-# stays cheap as the consumed-history grows.
-sa.Index(
- "idx_ws_attachments_reserved_at",
- workstream_attachments.c.reserved_at,
- sqlite_where=workstream_attachments.c.reserved_at.is_not(None),
- postgresql_where=workstream_attachments.c.reserved_at.is_not(None),
-)
# ---------------------------------------------------------------------------
# Skill versions — version history for skills
diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py
index 84d0dc2e..0cd03379 100644
--- a/turnstone/core/storage/_sqlite.py
+++ b/turnstone/core/storage/_sqlite.py
@@ -3,16 +3,18 @@
from __future__ import annotations
import contextlib
+import json
import queue
import threading
import time
+from collections import Counter
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any
import sqlalchemy as sa
if TYPE_CHECKING:
- from collections.abc import Iterable, Iterator
+ from collections.abc import Iterable, Iterator, Sequence
from turnstone.core.storage._notify import Notify, NotifyStream
@@ -106,12 +108,18 @@ from turnstone.core.storage._utils import (
from turnstone.core.storage._utils import (
VERDICT_MUTABLE as _VERDICT_MUTABLE,
)
+from turnstone.core.storage._utils import (
+ build_attachments_by_msg as _build_attachments_by_msg,
+)
from turnstone.core.storage._utils import (
escape_like as _escape_like,
)
from turnstone.core.storage._utils import (
normalize_search_terms as _normalize_search_terms,
)
+from turnstone.core.storage._utils import (
+ parse_attachment_refs as _parse_attachment_refs,
+)
from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text
from turnstone.core.storage._utils import (
reconstruct_messages as _reconstruct_messages,
@@ -422,24 +430,30 @@ class SQLiteBackend:
def load_messages(
self, ws_id: str, *, limit: int | None = None, repair: bool = True
) -> list[dict[str, Any]]:
+ # The trailing ``attachments`` column carries the per-row
+ # content-addressed ref-list; it is split off below to resolve blobs
+ # and is NOT part of the positional tuple ``reconstruct_messages``
+ # unpacks (id..is_error).
+ _cols = (
+ conversations.c.id,
+ conversations.c.role,
+ conversations.c.content,
+ conversations.c.tool_name,
+ conversations.c.tool_call_id,
+ conversations.c.provider_data,
+ conversations.c.tool_calls,
+ conversations.c._source,
+ conversations.c.event_id,
+ conversations.c.is_error,
+ conversations.c.attachments,
+ )
with self._conn() as conn:
if limit is not None and limit > 0:
# Tail-N: fetch the last `limit` rows via DESC + LIMIT
# then reverse so the reconstructed output stays in
# chronological order. Bounds memory on long histories.
rows = conn.execute(
- sa.select(
- conversations.c.id,
- conversations.c.role,
- conversations.c.content,
- conversations.c.tool_name,
- conversations.c.tool_call_id,
- conversations.c.provider_data,
- conversations.c.tool_calls,
- conversations.c._source,
- conversations.c.event_id,
- conversations.c.is_error,
- )
+ sa.select(*_cols)
.where(conversations.c.ws_id == ws_id)
.order_by(conversations.c.id.desc())
.limit(limit)
@@ -447,31 +461,37 @@ class SQLiteBackend:
rows = list(reversed(rows))
else:
rows = conn.execute(
- sa.select(
- conversations.c.id,
- conversations.c.role,
- conversations.c.content,
- conversations.c.tool_name,
- conversations.c.tool_call_id,
- conversations.c.provider_data,
- conversations.c.tool_calls,
- conversations.c._source,
- conversations.c.event_id,
- conversations.c.is_error,
- )
+ sa.select(*_cols)
.where(conversations.c.ws_id == ws_id)
.order_by(conversations.c.id)
).fetchall()
- # Bound the attachment scan to the fetched message ids when
- # tail-N was requested — otherwise the attachments query
- # still scans every row for the workstream and partially
- # defeats the conversations-table LIMIT.
- message_ids: list[int] | None = None
- if limit is not None and limit > 0:
- message_ids = [r[0] for r in rows]
- attachments = self.load_attachments_for_messages(ws_id, message_ids=message_ids)
- return _reconstruct_messages(list(rows), ws_id, attachments or None, repair=repair)
+ attachments = self._resolve_row_attachments(rows)
+ # Strip the trailing ref-list column so the tuple shape stays exactly
+ # what ``reconstruct_messages`` expects (id..is_error).
+ msg_rows = [tuple(r)[:10] for r in rows]
+ return _reconstruct_messages(msg_rows, ws_id, attachments or None, repair=repair)
+
+ def _resolve_row_attachments(self, rows: Sequence[Any]) -> dict[int, list[dict[str, Any]]]:
+ """Build the ``reconstruct_messages`` attachment map from row ref-lists.
+
+ Each row's trailing ``attachments`` column (last element) is the
+ content-addressed ref-list; collect every referenced id, bulk-fetch
+ the blobs in one query, and group them back per row id in ref-list
+ order. No referenced ids → no query.
+ """
+ attachment_refs: dict[int, list[str]] = {}
+ all_ids: set[str] = set()
+ for r in rows:
+ ids = _parse_attachment_refs(r[10])
+ if ids:
+ attachment_refs[r[0]] = ids
+ all_ids.update(ids)
+ if not all_ids:
+ return {}
+ blobs = self.get_attachments(list(all_ids))
+ rows_by_id = {str(b["attachment_id"]): b for b in blobs}
+ return _build_attachments_by_msg(attachment_refs, rows_by_id)
def get_max_event_id(self, ws_id: str) -> int | None:
with self._conn() as conn:
@@ -495,16 +515,23 @@ class SQLiteBackend:
if cutoff_row is None:
return 0 # nothing to delete
cutoff_id = cutoff_row[0]
- # Cascade-delete attachments linked to doomed messages so
- # rewind/retry flows don't leak orphan BLOBs.
- conn.execute(
- sa.delete(workstream_attachments).where(
+ # Refcount GC: read the doomed rows' content-addressed ref-lists,
+ # decrement each blob's refcount once per reference, and prune
+ # blobs that hit 0 — so a deduped blob still referenced by a kept
+ # turn survives. Replaces the old message_id-cascade delete.
+ doomed = conn.execute(
+ sa.select(conversations.c.attachments).where(
sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.message_id >= cutoff_id,
+ conversations.c.ws_id == ws_id,
+ conversations.c.id >= cutoff_id,
+ conversations.c.attachments.is_not(None),
)
)
- )
+ ).fetchall()
+ doomed_ids: list[str] = []
+ for (refs,) in doomed:
+ doomed_ids.extend(_parse_attachment_refs(refs))
+ self._release_attachment_refs(conn, doomed_ids)
# Remove FTS5 entries first (external content table doesn't auto-sync)
if self._fts5_available:
try:
@@ -962,9 +989,22 @@ class SQLiteBackend:
def delete_workstream(self, ws_id: str) -> bool:
with self._conn() as conn:
- conn.execute(
- sa.delete(workstream_attachments).where(workstream_attachments.c.ws_id == ws_id)
- )
+ # Refcount GC over every referenced blob (content-addressed ids are
+ # global, so a deduped blob may be shared with another workstream —
+ # decrement, don't blanket-delete by ws_id). Blobs that hit 0 are
+ # pruned; any still referenced elsewhere survive.
+ referenced = conn.execute(
+ sa.select(conversations.c.attachments).where(
+ sa.and_(
+ conversations.c.ws_id == ws_id,
+ conversations.c.attachments.is_not(None),
+ )
+ )
+ ).fetchall()
+ ref_ids: list[str] = []
+ for (refs,) in referenced:
+ ref_ids.extend(_parse_attachment_refs(refs))
+ self._release_attachment_refs(conn, ref_ids)
conn.execute(sa.delete(conversations).where(conversations.c.ws_id == ws_id))
conn.execute(sa.delete(workstream_config).where(workstream_config.c.ws_id == ws_id))
conn.execute(
@@ -985,7 +1025,7 @@ class SQLiteBackend:
conn.commit()
return result.rowcount > 0
- # -- Workstream attachments ------------------------------------------------
+ # -- Workstream attachments (content-addressed, refcounted) ----------------
def save_attachment(
self,
@@ -997,48 +1037,70 @@ class SQLiteBackend:
size_bytes: int,
kind: str,
content: bytes,
+ origin: str = "upload",
) -> None:
+ """Write a content-addressed blob (INSERT-OR-IGNORE) and ``refcount += 1``.
+
+ ``attachment_id`` is the content hash (the caller computes it). The
+ first reference writes the row at ``refcount = 1``; every subsequent
+ reference (a re-upload of identical bytes, or a second message
+ referencing the same blob) finds the PK present and only bumps the
+ count — so a stored blob is always referenced (born at ≥ 1) and dedupes
+ across messages / workstreams. Idempotent on the bytes, never on the
+ count.
+ """
+ from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
with self._conn() as conn:
+ # INSERT-OR-IGNORE the blob, then unconditionally bump the count.
+ # Splitting insert (ignore-on-conflict) from the increment keeps
+ # the +1 correct whether or not the row already existed.
+ stmt = sqlite_insert(workstream_attachments).values(
+ attachment_id=attachment_id,
+ ws_id=ws_id,
+ user_id=user_id,
+ filename=filename,
+ mime_type=mime_type,
+ size_bytes=size_bytes,
+ kind=kind,
+ content=content,
+ created=now,
+ refcount=0,
+ origin=origin,
+ )
+ conn.execute(stmt.on_conflict_do_nothing(index_elements=["attachment_id"]))
conn.execute(
- sa.insert(workstream_attachments),
- {
- "attachment_id": attachment_id,
- "ws_id": ws_id,
- "user_id": user_id,
- "filename": filename,
- "mime_type": mime_type,
- "size_bytes": size_bytes,
- "kind": kind,
- "content": content,
- "message_id": None,
- "created": now,
- },
+ sa.update(workstream_attachments)
+ .where(workstream_attachments.c.attachment_id == attachment_id)
+ .values(refcount=workstream_attachments.c.refcount + 1)
)
conn.commit()
- def list_pending_attachments(self, ws_id: str, user_id: str) -> list[dict[str, Any]]:
+ def set_message_attachments(
+ self, ws_id: str, message_id: int, attachment_ids: list[str]
+ ) -> None:
+ """Record a turn's ordered content-addressed ref-list on its row.
+
+ Writes the JSON id-list onto ``conversations.attachments`` for the
+ ``(ws_id, message_id)`` row — the sole message->blob link. Empty
+ input is a no-op (the column stays NULL). Scoped to ``ws_id`` as
+ defense-in-depth against a cross-ws message id.
+ """
+ if not attachment_ids or not message_id:
+ return
with self._conn() as conn:
- rows = conn.execute(
- sa.select(
- workstream_attachments.c.attachment_id,
- workstream_attachments.c.filename,
- workstream_attachments.c.mime_type,
- workstream_attachments.c.size_bytes,
- workstream_attachments.c.kind,
- workstream_attachments.c.created,
- )
+ conn.execute(
+ sa.update(conversations)
.where(
sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
+ conversations.c.id == message_id,
+ conversations.c.ws_id == ws_id,
)
)
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- return [dict(r._mapping) for r in rows]
+ .values(attachments=json.dumps(list(attachment_ids)))
+ )
+ conn.commit()
def get_attachments(self, attachment_ids: list[str]) -> list[dict[str, Any]]:
if not attachment_ids:
@@ -1051,24 +1113,6 @@ class SQLiteBackend:
).fetchall()
return [dict(r._mapping) for r in rows]
- def get_pending_attachments_with_content(
- self, ws_id: str, user_id: str
- ) -> list[dict[str, Any]]:
- with self._conn() as conn:
- rows = conn.execute(
- sa.select(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
- )
- )
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- return [dict(r._mapping) for r in rows]
-
def get_attachment(self, attachment_id: str) -> dict[str, Any] | None:
with self._conn() as conn:
row = conn.execute(
@@ -1078,163 +1122,59 @@ class SQLiteBackend:
).fetchone()
return dict(row._mapping) if row else None
- def delete_attachment(self, attachment_id: str, ws_id: str, user_id: str) -> bool:
+ def attachment_referenced_in_ws(self, attachment_id: str, ws_id: str) -> bool:
+ """True iff some conversations row in ``ws_id`` references ``attachment_id``.
+
+ The committed-attachment ownership gate: the ``ws_id``/``user_id``
+ scope columns are gone, so a ``get_content`` for a committed blob is
+ authorised by proving the requester (already gated to own ``ws_id``)
+ has a turn in that workstream whose ref-list names the id. Uses a
+ JSON-array substring match on the ``attachments`` column —
+ content-addressed ids are 64-char sha256 hex, so a quoted-id substring
+ cannot collide with another id.
+ """
+ needle = f'%"{attachment_id}"%'
with self._conn() as conn:
- # Only pending (unreserved, unconsumed) attachments may be
- # deleted. Reserved ones are soft-locked to a queued send.
- result = conn.execute(
- sa.delete(workstream_attachments).where(
+ row = conn.execute(
+ sa.select(conversations.c.id)
+ .where(
sa.and_(
- workstream_attachments.c.attachment_id == attachment_id,
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
+ conversations.c.ws_id == ws_id,
+ conversations.c.attachments.is_not(None),
+ conversations.c.attachments.like(needle),
)
)
- )
- conn.commit()
- return result.rowcount > 0
+ .limit(1)
+ ).fetchone()
+ return row is not None
- def mark_attachments_consumed(
- self,
- attachment_ids: list[str],
- message_id: int,
- ws_id: str,
- user_id: str,
- reserved_for_msg_id: str | None = None,
- ) -> None:
+ @staticmethod
+ def _release_attachment_refs(conn: Any, attachment_ids: list[str]) -> None:
+ """Decrement refcount once per id and prune blobs that reach 0.
+
+ Caller holds the connection / transaction. Counts duplicate ids in
+ the input (a turn references an id once, but a batch may span several
+ turns that each reference the same deduped blob), so the decrement
+ matches the number of references actually being removed.
+ """
if not attachment_ids:
return
- predicate = sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
+ counts = Counter(attachment_ids)
+ for aid, n in counts.items():
+ conn.execute(
+ sa.update(workstream_attachments)
+ .where(workstream_attachments.c.attachment_id == aid)
+ .values(refcount=workstream_attachments.c.refcount - n)
+ )
+ # Prune any blob whose count fell to (or below) 0.
+ conn.execute(
+ sa.delete(workstream_attachments).where(
+ sa.and_(
+ workstream_attachments.c.attachment_id.in_(list(counts)),
+ workstream_attachments.c.refcount <= 0,
+ )
+ )
)
- if reserved_for_msg_id is not None:
- predicate = sa.and_(
- predicate,
- workstream_attachments.c.reserved_for_msg_id == reserved_for_msg_id,
- )
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(predicate)
- .values(
- message_id=message_id,
- reserved_for_msg_id=None,
- reserved_at=None,
- )
- )
- conn.commit()
-
- def reserve_attachments(
- self,
- attachment_ids: list[str],
- queue_msg_id: str,
- ws_id: str,
- user_id: str,
- ) -> list[str]:
- if not attachment_ids or not queue_msg_id:
- return []
- now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_for_msg_id.is_(None),
- )
- )
- .values(reserved_for_msg_id=queue_msg_id, reserved_at=now)
- )
- # Echo back which ids are now reserved for this msg id (race-
- # safe confirmation for the caller).
- rows = conn.execute(
- sa.select(workstream_attachments.c.attachment_id).where(
- sa.and_(
- workstream_attachments.c.attachment_id.in_(attachment_ids),
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.reserved_for_msg_id == queue_msg_id,
- )
- )
- ).fetchall()
- conn.commit()
- return [r[0] for r in rows]
-
- def unreserve_attachments(self, queue_msg_id: str, ws_id: str, user_id: str) -> None:
- if not queue_msg_id:
- return
- with self._conn() as conn:
- conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.user_id == user_id,
- workstream_attachments.c.reserved_for_msg_id == queue_msg_id,
- )
- )
- .values(reserved_for_msg_id=None, reserved_at=None)
- )
- conn.commit()
-
- def sweep_orphan_reservations(self, older_than_seconds: int) -> int:
- if older_than_seconds <= 0:
- return 0
- cutoff = (datetime.now(UTC) - timedelta(seconds=older_than_seconds)).strftime(
- "%Y-%m-%dT%H:%M:%S"
- )
- with self._conn() as conn:
- result = conn.execute(
- sa.update(workstream_attachments)
- .where(
- sa.and_(
- workstream_attachments.c.reserved_for_msg_id.is_not(None),
- workstream_attachments.c.message_id.is_(None),
- workstream_attachments.c.reserved_at.is_not(None),
- workstream_attachments.c.reserved_at < cutoff,
- )
- )
- .values(reserved_for_msg_id=None, reserved_at=None)
- )
- conn.commit()
- return int(result.rowcount or 0)
-
- def load_attachments_for_messages(
- self,
- ws_id: str,
- *,
- message_ids: list[int] | None = None,
- ) -> dict[int, list[dict[str, Any]]]:
- with self._conn() as conn:
- where_clauses = [
- workstream_attachments.c.ws_id == ws_id,
- workstream_attachments.c.message_id.is_not(None),
- ]
- if message_ids is not None:
- # Empty list → no matches; guard against an implicit
- # all-rows scan from a would-be empty IN clause.
- if not message_ids:
- return {}
- where_clauses.append(workstream_attachments.c.message_id.in_(message_ids))
- rows = conn.execute(
- sa.select(workstream_attachments)
- .where(sa.and_(*where_clauses))
- .order_by(workstream_attachments.c.created)
- ).fetchall()
- grouped: dict[int, list[dict[str, Any]]] = {}
- for r in rows:
- row = dict(r._mapping)
- mid = row["message_id"]
- grouped.setdefault(mid, []).append(row)
- return grouped
def list_workstreams(
self,
diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py
index 7ccdf0fb..fc5a9ab2 100644
--- a/turnstone/core/storage/_utils.py
+++ b/turnstone/core/storage/_utils.py
@@ -156,6 +156,80 @@ def _attachment_to_content_part(att: dict[str, Any]) -> dict[str, Any] | None:
return None
+def parse_attachment_refs(raw: str | None) -> list[str]:
+ """Decode a ``conversations.attachments`` ref-list column into id strings.
+
+ The column stores a JSON array of content-addressed ``attachment_id``s in
+ turn order (NULL / empty for turns with no attachments). Malformed or
+ non-list JSON decodes to an empty list (defensive — a corrupt column must
+ never crash a history load); non-string elements are dropped.
+ """
+ if not raw:
+ return []
+ try:
+ parsed = json.loads(raw)
+ except (json.JSONDecodeError, TypeError):
+ return []
+ if not isinstance(parsed, list):
+ return []
+ return [str(x) for x in parsed if isinstance(x, str) and x]
+
+
+def build_attachments_by_msg(
+ attachment_refs: dict[int, list[str]],
+ rows_by_id: dict[str, dict[str, Any]],
+) -> dict[int, list[dict[str, Any]]]:
+ """Assemble the ``reconstruct_messages`` attachment map from ref-lists.
+
+ ``attachment_refs`` maps a conversations row id to its ordered list of
+ content-addressed ids (from :func:`parse_attachment_refs`); ``rows_by_id``
+ maps an attachment id to its resolved blob row (incl. ``content`` bytes).
+ Returns ``{row_id: [att_row, ...]}`` preserving ref-list order, skipping
+ ids whose blob is missing (pruned / never written). Empty lists are
+ omitted so the caller can pass ``result or None`` unchanged.
+ """
+ grouped: dict[int, list[dict[str, Any]]] = {}
+ for mid, ids in attachment_refs.items():
+ resolved = [rows_by_id[aid] for aid in ids if aid in rows_by_id]
+ if resolved:
+ grouped[mid] = resolved
+ return grouped
+
+
+def _reconstruct_attachment_parts(
+ attachments_by_msg: dict[int, list[dict[str, Any]]] | None,
+ row_id: int | None,
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+ """Build ``(content_parts, attachments_meta)`` for a row from its ref-list.
+
+ ``attachments_by_msg`` maps a conversations row id to the ordered list of
+ content-addressed attachment rows referenced by that row's
+ ``attachments`` column (resolved to include ``content`` bytes). Returns
+ the OpenAI-style content parts (image_url / document, in ref-list order)
+ and the display-oriented ``_attachments_meta`` siblings (kind / filename /
+ mime_type) — the latter is tracked even when a part can't be reconstructed
+ so history replay keeps filenames available (e.g. image pills). Shared by
+ the user- and tool-row reconstruction so both surfaces stay byte-identical
+ in part shape.
+ """
+ parts: list[dict[str, Any]] = []
+ meta: list[dict[str, Any]] = []
+ if not attachments_by_msg or row_id is None:
+ return parts, meta
+ for att in attachments_by_msg.get(row_id, []):
+ part = _attachment_to_content_part(att)
+ if part is not None:
+ parts.append(part)
+ meta.append(
+ {
+ "kind": str(att.get("kind") or ""),
+ "filename": str(att.get("filename") or ""),
+ "mime_type": str(att.get("mime_type") or ""),
+ }
+ )
+ return parts, meta
+
+
# ---------------------------------------------------------------------------
# Search-term normalization
# ---------------------------------------------------------------------------
@@ -453,9 +527,13 @@ def reconstruct_messages(
cursor) is surfaced as the ``_event_id`` side-channel; legacy 9-tuple
fixtures omit it (handled by the defensive unpack below).
- When ``attachments_by_msg`` is provided, any user row whose id has
- attachments is rebuilt with multipart list content (text +
- image_url/document parts).
+ When ``attachments_by_msg`` is provided (keyed by row id, each value an
+ ordered list of content-addressed attachment rows resolved from the
+ ``conversations.attachments`` ref-list), any ``user`` *or* ``tool`` row
+ whose id has attachments is rebuilt with multipart list content (text +
+ image_url/document parts). Tool rows carry persisted vision output
+ (``read_file`` on an image) this way — they would otherwise reload as the
+ flattened text alone.
When ``repair`` is True (default) the result is post-processed to
produce a wire-shape valid for an LLM round-trip: the trailing
@@ -494,23 +572,7 @@ def reconstruct_messages(
is_error = bool(row[9]) if len(row) > 9 else False
if role == "user":
- parts: list[dict[str, Any]] = []
- meta: list[dict[str, Any]] = []
- if attachments_by_msg and row_id is not None:
- for att in attachments_by_msg.get(row_id, []):
- part = _attachment_to_content_part(att)
- if part is not None:
- parts.append(part)
- # Track display-oriented metadata even when a part
- # itself can't be reconstructed — keeps filenames
- # available for history replay (e.g. image pills).
- meta.append(
- {
- "kind": str(att.get("kind") or ""),
- "filename": str(att.get("filename") or ""),
- "mime_type": str(att.get("mime_type") or ""),
- }
- )
+ parts, meta = _reconstruct_attachment_parts(attachments_by_msg, row_id)
if parts:
user_content: list[dict[str, Any]] = [{"type": "text", "text": content or ""}]
user_content.extend(parts)
@@ -547,10 +609,22 @@ def reconstruct_messages(
messages.append(msg)
elif role == "tool":
+ # Tool rows can carry persisted vision output (``read_file`` on an
+ # image): the live tool message's content was a multipart list and
+ # the image bytes were written content-addressed + referenced on
+ # this row, while the row's text column holds only the flattened
+ # text. Rebuild the multipart list on reload so the image survives;
+ # text-only tool rows stay plain strings (the common case).
+ tparts, _tmeta = _reconstruct_attachment_parts(attachments_by_msg, row_id)
+ tool_content: str | list[dict[str, Any]]
+ if tparts:
+ tool_content = [{"type": "text", "text": content or ""}, *tparts]
+ else:
+ tool_content = content or ""
tmsg: dict[str, Any] = {
"role": "tool",
"tool_call_id": tc_id or "",
- "content": content or "",
+ "content": tool_content,
}
if is_error:
tmsg["is_error"] = True
diff --git a/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py b/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py
index dbbbad80..b0cb516d 100644
--- a/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py
+++ b/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py
@@ -33,6 +33,15 @@ generating provider as the ``{producer, blocks}`` envelope (producer inferred fr
types — see ``_infer_producer``, which must match the live save's ``provider_name``
values), so the lowering layer can replay the native lane verbatim only to its producer.
+Finally it performs the **attachment content-addressing cutover**: after adding
+``conversations.attachments`` (the ref-list) and ``workstream_attachments.refcount`` /
+``origin``, it re-keys every legacy *consumed* attachment row to its content hash
+(``sha256(content)``), dedups identical bytes into one refcounted blob, builds each
+message's ``attachments`` ref-list from the old ``message_id`` link, and then drops the
+retired upload-lifecycle columns ``message_id`` / ``reserved_for_msg_id`` /
+``reserved_at`` (and their indexes). Pending (un-consumed) legacy rows are dropped —
+pending uploads now live in the per-node in-memory buffer, not in storage.
+
``downgrade()`` re-adds the (empty) ``_reminders`` column so the schema matches
the 059 state, but does NOT reverse the envelope un-wrap — that is lossy (the
advisory blocks are discarded), so the original wrapped rows cannot be
@@ -45,6 +54,7 @@ Create Date: 2026-06-01
from __future__ import annotations
+import hashlib
import json
from typing import Any
@@ -253,8 +263,9 @@ def upgrade() -> None:
batch_op.add_column(
sa.Column("is_error", sa.Boolean, nullable=False, server_default=sa.false())
)
- # Content-addressed attachment ref-list (canonical-trajectory cut); the cutover
- # fills it and retires the message_id/reserved_* link.
+ # Content-addressed attachment ref-list (canonical-trajectory cut); the
+ # backfill below (step 3) fills it from the legacy message_id link and
+ # then drops message_id/reserved_* (step 4).
batch_op.add_column(sa.Column("attachments", sa.Text, nullable=True))
with op.batch_alter_table("workstream_attachments") as batch_op:
batch_op.add_column(
@@ -264,16 +275,159 @@ def upgrade() -> None:
sa.Column("origin", sa.Text, nullable=False, server_default=sa.text("'upload'"))
)
+ # (3) Backfill the content-addressed model from the legacy message_id link,
+ # then drop the retired lifecycle columns. Must run AFTER the additive
+ # columns above exist (refcount / origin / attachments) and BEFORE the
+ # drop below (it reads message_id).
+ _backfill_content_addressed_attachments(bind)
+
+ # (4) Drop the retired upload-lifecycle columns + their indexes. The
+ # content-addressed model keys blobs by content hash and links them via
+ # the conversations.attachments ref-list, so message_id /
+ # reserved_for_msg_id / reserved_at (and the indexes over them) are dead.
+ # Drop the dependent indexes FIRST on both dialects: SQLite's
+ # ``batch_alter_table`` rebuilds the table from the reflected schema and
+ # would otherwise try to re-create these indexes against the
+ # now-missing columns; PostgreSQL needs them gone before the columns.
+ for idx in (
+ "idx_ws_attachments_pending",
+ "idx_ws_attachments_message",
+ "idx_ws_attachments_reserved",
+ "idx_ws_attachments_reserved_at",
+ ):
+ op.execute(sa.text(f"DROP INDEX IF EXISTS {idx}"))
+ with op.batch_alter_table("workstream_attachments") as batch_op:
+ batch_op.drop_column("message_id")
+ batch_op.drop_column("reserved_for_msg_id")
+ batch_op.drop_column("reserved_at")
+
+
+def _backfill_content_addressed_attachments(bind: sa.engine.Connection) -> None:
+ """Re-key legacy consumed attachments to their content hash + build ref-lists.
+
+ Legacy rows linked an attachment to a message via
+ ``workstream_attachments.message_id``. The content-addressed model keys a
+ blob by ``sha256(content)`` and links it via the ordered
+ ``conversations.attachments`` ref-list. For every *consumed* legacy row
+ (``message_id IS NOT NULL``):
+
+ * compute the content hash and dedup — identical bytes collapse to one row
+ whose PK is re-keyed to the hash; duplicate legacy rows are deleted;
+ * set ``refcount`` = the number of distinct messages referencing that
+ content, and ``origin = 'upload'``;
+ * build each referencing message's ``conversations.attachments`` as the
+ ordered list of content hashes (legacy per-message order preserved by the
+ attachment row's ``created`` then ``attachment_id``).
+
+ Pending (un-consumed) legacy rows (``message_id IS NULL``) are dropped: they
+ were transient upload state and the content-addressed model holds no pending
+ blobs in storage (they live in the per-node buffer now).
+ """
+ wa = sa.table(
+ "workstream_attachments",
+ sa.column("attachment_id", sa.Text),
+ sa.column("message_id", sa.Integer),
+ sa.column("content", sa.LargeBinary),
+ sa.column("created", sa.Text),
+ sa.column("refcount", sa.Integer),
+ sa.column("origin", sa.Text),
+ )
+ conversations = sa.table(
+ "conversations",
+ sa.column("id", sa.Integer),
+ sa.column("attachments", sa.Text),
+ )
+
+ # Read every consumed legacy row in (message_id, created, attachment_id)
+ # order so each message's ref-list preserves the original attachment order.
+ rows = bind.execute(
+ sa.select(wa.c.attachment_id, wa.c.message_id, wa.c.content)
+ .where(wa.c.message_id.is_not(None))
+ .order_by(wa.c.message_id, wa.c.created, wa.c.attachment_id)
+ ).fetchall()
+
+ # new_id (content hash) -> canonical old id kept as that blob's row.
+ canonical_old_id: dict[str, str] = {}
+ # new_id -> set of distinct message ids referencing it (refcount source).
+ refcounting: dict[str, set[int]] = {}
+ # message_id -> ordered list of new_ids (de-duped within the message).
+ per_message: dict[int, list[str]] = {}
+ # old ids to delete (duplicates that collapsed into a canonical row).
+ drop_old_ids: list[str] = []
+
+ for old_id, message_id, content in rows:
+ raw = content if isinstance(content, (bytes, bytearray)) else b""
+ new_id = hashlib.sha256(bytes(raw)).hexdigest()
+ if new_id not in canonical_old_id:
+ canonical_old_id[new_id] = old_id
+ refcounting[new_id] = set()
+ elif old_id != canonical_old_id[new_id]:
+ # A distinct legacy row carrying identical bytes — collapse it.
+ drop_old_ids.append(old_id)
+ refcounting[new_id].add(int(message_id))
+ bucket = per_message.setdefault(int(message_id), [])
+ if new_id not in bucket:
+ bucket.append(new_id)
+
+ # Re-key each canonical row's PK to its content hash and set refcount/origin.
+ # Re-key first (while the duplicates still hold their old PKs), then delete
+ # the duplicates, so a re-key can't collide with a not-yet-deleted dup.
+ for new_id, old_id in canonical_old_id.items():
+ bind.execute(
+ sa.update(wa)
+ .where(wa.c.attachment_id == old_id)
+ .values(
+ attachment_id=new_id,
+ refcount=len(refcounting[new_id]),
+ origin="upload",
+ )
+ )
+ for old_id in drop_old_ids:
+ bind.execute(sa.delete(wa).where(wa.c.attachment_id == old_id))
+
+ # Drop any remaining pending (un-consumed) legacy rows — no storage home.
+ bind.execute(sa.delete(wa).where(wa.c.message_id.is_(None)))
+
+ # Write each message's content-addressed ref-list.
+ for message_id, new_ids in per_message.items():
+ bind.execute(
+ sa.update(conversations)
+ .where(conversations.c.id == message_id)
+ .values(attachments=json.dumps(new_ids))
+ )
+
def downgrade() -> None:
# Re-add the (empty) column so the schema matches the 059 state. The
# envelope un-wrap (step 1) is NOT reversed — it discards the advisory
# blocks, so the original wrapped rows cannot be reconstructed; the
- # re-added column is therefore always NULL.
+ # re-added column is therefore always NULL. The content-addressing
+ # backfill (step 3) is likewise NOT reversed: the retired columns are
+ # re-added empty (the old message_id links / reservation tokens cannot be
+ # reconstructed from the content-addressed ref-list).
with op.batch_alter_table("conversations") as batch_op:
batch_op.add_column(sa.Column("_reminders", sa.Text, nullable=True))
batch_op.drop_column("is_error")
batch_op.drop_column("attachments")
with op.batch_alter_table("workstream_attachments") as batch_op:
+ batch_op.add_column(sa.Column("message_id", sa.Integer, nullable=True))
+ batch_op.add_column(sa.Column("reserved_for_msg_id", sa.Text, nullable=True))
+ batch_op.add_column(sa.Column("reserved_at", sa.Text, nullable=True))
batch_op.drop_column("refcount")
batch_op.drop_column("origin")
+ # Re-create the indexes over the re-added columns to match the 059 schema.
+ op.create_index(
+ "idx_ws_attachments_pending",
+ "workstream_attachments",
+ ["ws_id", "user_id", "message_id"],
+ )
+ op.create_index(
+ "idx_ws_attachments_message",
+ "workstream_attachments",
+ ["message_id"],
+ )
+ op.create_index(
+ "idx_ws_attachments_reserved",
+ "workstream_attachments",
+ ["ws_id", "user_id", "reserved_for_msg_id"],
+ )
diff --git a/turnstone/server.py b/turnstone/server.py
index 27599e1d..0bc2a68c 100644
--- a/turnstone/server.py
+++ b/turnstone/server.py
@@ -118,14 +118,6 @@ _VALID_WS_ID = re.compile(r"^[0-9a-f]{32}$")
# WebUI — implements SessionUI for browser-based interaction
# ---------------------------------------------------------------------------
-# Orphan-attachment-reservation sweep cadence. Threshold is measured
-# against the storage layer's `reserved_at` column (time the row last
-# transitioned into reserved state, NOT upload time), so a 1-hour cap
-# is safely longer than any realistic single send without risking the
-# unreservation of attachments uploaded long ago but reserved fresh.
-_ORPHAN_SWEEP_INTERVAL_S = 30 * 60
-_ORPHAN_SWEEP_THRESHOLD_S = 1 * 3600
-
class WebUI(SessionUIBase):
"""Browser-based UI using SSE for streaming and HTTP POST for actions.
@@ -2050,14 +2042,16 @@ async def _interactive_create_post_install(
initial_message = body.get("initial_message", "").strip()
if initial_message and ws.session is not None:
from turnstone.core.attachments import (
- reserve_and_resolve_attachments as _reserve_and_resolve,
+ resolve_staged_attachments as _resolve_staged,
)
session = ws.session
send_id = uuid.uuid4().hex
resolved_atts: list[Any] = []
if attachment_ids:
- resolved_atts, _ord, _drop = _reserve_and_resolve(attachment_ids, send_id, ws.id, uid)
+ # Resolve (peek) the staged uploads; the committing send drains
+ # them from the buffer. No reservation to release on failure.
+ resolved_atts, _ord, _drop = _resolve_staged(attachment_ids, ws.id, uid)
def _run_initial() -> None:
try:
@@ -2067,13 +2061,6 @@ async def _interactive_create_post_install(
send_id=send_id if resolved_atts else None,
)
except (Exception, GenerationCancelled):
- if attachment_ids:
- from turnstone.core.memory import (
- unreserve_attachments as _unreserve,
- )
-
- with contextlib.suppress(Exception):
- _unreserve(send_id, ws.id, uid)
if isinstance(ws.ui, WebUI):
ws.ui.on_stream_end()
ws.ui.on_state_change("idle")
@@ -3472,36 +3459,9 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
if state_writer is not None:
state_writer.start()
- # Sweep stale attachment reservations left over from process crashes
- # between reserve_attachments and consume/unreserve. Run once at
- # startup (catches anything orphaned by the previous process), then
- # periodically as defense-in-depth.
- from turnstone.core.memory import sweep_orphan_reservations as _sweep_orphans
-
- try:
- n = await asyncio.to_thread(_sweep_orphans, _ORPHAN_SWEEP_THRESHOLD_S)
- if n:
- log.info("attachments.orphan_sweep.startup", swept=n)
- except Exception:
- log.warning("attachments.orphan_sweep.startup_failed", exc_info=True)
-
- _orphan_sweep_stop = asyncio.Event()
-
- async def _orphan_sweep_loop() -> None:
- while not _orphan_sweep_stop.is_set():
- try:
- await asyncio.wait_for(_orphan_sweep_stop.wait(), timeout=_ORPHAN_SWEEP_INTERVAL_S)
- return # stop event fired
- except TimeoutError:
- pass
- try:
- n = await asyncio.to_thread(_sweep_orphans, _ORPHAN_SWEEP_THRESHOLD_S)
- if n:
- log.info("attachments.orphan_sweep.periodic", swept=n)
- except Exception:
- log.warning("attachments.orphan_sweep.periodic_failed", exc_info=True)
-
- _orphan_sweep_task = asyncio.create_task(_orphan_sweep_loop())
+ # (The attachment orphan-reservation sweep is gone — pending uploads now
+ # live in the per-node in-memory buffer with its own TTL eviction, so
+ # there are no persisted reservations to reclaim.)
from turnstone.core.oidc import initialize_oidc_state
@@ -3660,10 +3620,6 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
state_writer = getattr(app.state, "state_writer", None)
if state_writer is not None:
await asyncio.to_thread(state_writer.shutdown)
- # Stop the orphan-reservation sweep loop
- _orphan_sweep_stop.set()
- with contextlib.suppress(asyncio.CancelledError, Exception):
- await _orphan_sweep_task
# health_registry is stateless (no background threads) — nothing to stop
if app.state.mcp_client:
app.state.mcp_client.shutdown()
@@ -3782,14 +3738,10 @@ def create_app(
from turnstone.core.attachments import (
sniff_image_mime as _sniff_image_mime,
)
- from turnstone.core.attachments import (
- upload_lock as _attachment_upload_lock,
- )
interactive_attachment_helpers = AttachmentUploadHelpers(
sniff_image_mime=_sniff_image_mime,
classify_text_attachment=_classify_text_attachment,
- upload_lock=_attachment_upload_lock,
)
from turnstone.core.memory import (
get_workstream_display_names as _get_ws_display_names,