feat(storage): content-addressed refcounted attachments + in-memory upload buffer

Replace the persisted pending/reserved/consumed upload lifecycle (and its orphan-sweep
and per-user cap) with a content-addressed, refcounted blob store fronted by the per-node
in-memory pending buffer:

- Upload stages bytes in the buffer (keyed by sha256); send-commit drains the referenced
  handles, writes each blob content-addressed (INSERT-OR-IGNORE then refcount += 1, so a
  stored blob is born referenced and dedupes across messages/workstreams), and records the
  ordered conversations.attachments ref-list — the sole message->blob link.
- reconstruct rebuilds inline image_url/document multipart content from the ref-list,
  role-agnostically (so tool-produced images via _exec_read_image now persist + rehydrate
  instead of being flattened to text and lost). Output shape unchanged.
- GC is reference counting: delete_messages_after / delete_workstream decrement once per
  reference and prune a blob at 0; a deduped blob shared with a kept turn (or another ws)
  survives.
- get_content for a committed blob is gated by reference-ownership (the requester owns a
  turn in the ws whose ref-list names the id), replacing the dropped ws_id/user_id scope.
- Migration 060 re-keys legacy consumed attachments to their content hash, dedups, sets
  refcounts, writes the ref-lists, and drops message_id/reserved_*; pending legacy rows are
  dropped (pending now lives only in the buffer).

Both backends symmetric; the reservation methods, cap, and orphan-sweep are removed across
storage/facade/protocol/endpoints/coordinator. Wire harness byte-identical; full suite green.
This commit is contained in:
Patrick Buckley
2026-06-02 23:50:55 -07:00
parent 2e785fbdcf
commit 98cee4d20c
22 changed files with 1885 additions and 2181 deletions
+41 -40
View File
@@ -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())
+1 -1
View File
@@ -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)
+202 -1
View File
@@ -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
# ``<tool_output>`` block, then ``"\n".join`` with a part that itself begins
# with ``\n<system-reminder>`` — yielding the ``</tool_output>\n\n<system-
@@ -166,7 +194,9 @@ class TestMigration060:
finally:
engine.dispose()
def test_content_addressed_attachment_columns_added(self, tmp_path: Path) -> 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()
+60
View File
@@ -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"
+42 -364
View File
@@ -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
+138 -88
View File
@@ -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
+57 -43
View File
@@ -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.
+294 -414
View File
@@ -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
+23 -17
View File
@@ -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:
+19 -49
View File
@@ -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,
+6 -14
View File
@@ -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,
+101 -177
View File
@@ -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]:
+23 -108
View File
@@ -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.
+120 -26
View File
@@ -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:<mime>;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
+97 -141
View File
@@ -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"})
+169 -236
View File
@@ -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,
+35 -103
View File
@@ -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.
"""
...
+16 -38
View File
@@ -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 = <queue-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
+182 -242
View File
@@ -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,
+95 -21
View File
@@ -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
@@ -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"],
)
+7 -55
View File
@@ -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,