mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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())
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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"})
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user