wehub-resource-sync
227 lines
7.5 KiB
Python
227 lines
7.5 KiB
Python
"""Forking a session carries its attachment/artifact material to the child.
|
|
|
|
A fork copies transcript rows, but the artifact and attachment material stores are
|
|
keyed by session id. Without copying the material a forked conversation references
|
|
generated images/files and uploaded attachments that resolve to an empty child bucket
|
|
and fail to preview or replay. These tests pin the copy behavior at the storage layer
|
|
and end-to-end through ``SessionManager.branch``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from opensquilla.artifacts import ArtifactStore
|
|
from opensquilla.attachment_refs import (
|
|
copy_transcript_material,
|
|
make_attachment_ref,
|
|
read_attachment_ref_bytes,
|
|
transcript_material_path,
|
|
write_transcript_material,
|
|
)
|
|
from opensquilla.session.manager import SessionManager
|
|
from opensquilla.session.models import TranscriptEntry
|
|
from opensquilla.session.storage import SessionStorage
|
|
|
|
|
|
def _png_bytes() -> bytes:
|
|
from PIL import Image
|
|
|
|
out = io.BytesIO()
|
|
Image.new("RGB", (8, 8), color="red").save(out, format="PNG")
|
|
return out.getvalue()
|
|
|
|
|
|
def test_copy_transcript_material_to_child(tmp_path: Path) -> None:
|
|
sha, _path, wrote = write_transcript_material(
|
|
media_root=tmp_path, session_id="parent-1", payload=b"attachment payload"
|
|
)
|
|
assert wrote is True
|
|
|
|
copied = copy_transcript_material(
|
|
media_root=tmp_path, source_session_id="parent-1", target_session_id="child-1"
|
|
)
|
|
assert copied == 1
|
|
|
|
child_blob = transcript_material_path(tmp_path, "child-1", sha)
|
|
assert child_blob.exists()
|
|
assert child_blob.read_bytes() == b"attachment payload"
|
|
|
|
# A child-scoped ref now reads its own copy (replay resolves by current session).
|
|
child_ref = make_attachment_ref(
|
|
sha256=sha,
|
|
name="f.bin",
|
|
mime="application/octet-stream",
|
|
size=len(b"attachment payload"),
|
|
session_id="child-1",
|
|
source="transcript",
|
|
)
|
|
assert read_attachment_ref_bytes(child_ref, media_root=tmp_path) == b"attachment payload"
|
|
|
|
# Idempotent: re-copying materializes nothing new.
|
|
assert (
|
|
copy_transcript_material(
|
|
media_root=tmp_path, source_session_id="parent-1", target_session_id="child-1"
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
def test_copy_transcript_material_missing_source_is_noop(tmp_path: Path) -> None:
|
|
assert (
|
|
copy_transcript_material(
|
|
media_root=tmp_path, source_session_id="absent", target_session_id="child-1"
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def storage():
|
|
store = SessionStorage(":memory:")
|
|
await store.connect()
|
|
yield store
|
|
await store.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_copies_artifact_and_attachment_material(
|
|
storage: SessionStorage, tmp_path: Path
|
|
) -> None:
|
|
media_root = tmp_path / "media"
|
|
manager = SessionManager(storage, inject_time_prefix=False, media_root=media_root)
|
|
|
|
parent = await manager.create("agent:main:main")
|
|
|
|
artifact = ArtifactStore(media_root).publish_bytes(
|
|
_png_bytes(),
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
name="generated-image.png",
|
|
mime="image/png",
|
|
source="publish_artifact",
|
|
)
|
|
att_sha, _p, _w = write_transcript_material(
|
|
media_root=media_root, session_id=parent.session_id, payload=b"upload-bytes"
|
|
)
|
|
# A transcript entry must exist for fork_transcript to copy (and material to follow).
|
|
await storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
role="assistant",
|
|
content="here is your image",
|
|
token_count=4,
|
|
)
|
|
)
|
|
|
|
child = await manager.branch(
|
|
"agent:main:main", "agent:main:direct:u1", fork_transcript=True
|
|
)
|
|
assert child.forked_from_parent is True
|
|
|
|
# The forked child resolves the generated artifact under its own session id.
|
|
child_ref, child_path = ArtifactStore(media_root).resolve_for_download(
|
|
artifact.id, session_id=child.session_id
|
|
)
|
|
assert child_path.read_bytes() == _png_bytes()
|
|
assert child_ref.session_id == child.session_id
|
|
|
|
# And the uploaded attachment blob exists under the child's transcript store.
|
|
assert transcript_material_path(media_root, child.session_id, att_sha).exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_without_media_root_is_safe(storage: SessionStorage) -> None:
|
|
# No media_root configured: fork still succeeds, material copy is a no-op.
|
|
manager = SessionManager(storage, inject_time_prefix=False)
|
|
parent = await manager.create("agent:main:main")
|
|
await storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
role="assistant",
|
|
content="hi",
|
|
token_count=1,
|
|
)
|
|
)
|
|
child = await manager.branch(
|
|
"agent:main:main", "agent:main:direct:u2", fork_transcript=True
|
|
)
|
|
assert child.forked_from_parent is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_nested_fork_carries_material_each_generation(
|
|
storage: SessionStorage, tmp_path: Path
|
|
) -> None:
|
|
media_root = tmp_path / "media"
|
|
manager = SessionManager(storage, inject_time_prefix=False, media_root=media_root)
|
|
|
|
parent = await manager.create("agent:main:main")
|
|
artifact = ArtifactStore(media_root).publish_bytes(
|
|
_png_bytes(),
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
name="image.png",
|
|
mime="image/png",
|
|
source="publish_artifact",
|
|
)
|
|
await storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
role="assistant",
|
|
content="image",
|
|
token_count=2,
|
|
)
|
|
)
|
|
|
|
child = await manager.branch(
|
|
"agent:main:main", "agent:main:direct:child", fork_transcript=True
|
|
)
|
|
grandchild = await manager.branch(
|
|
"agent:main:direct:child", "agent:main:direct:grandchild", fork_transcript=True
|
|
)
|
|
|
|
# The artifact re-resolves under each generation's own session id.
|
|
for session_id in (child.session_id, grandchild.session_id):
|
|
_ref, path = ArtifactStore(media_root).resolve_for_download(
|
|
artifact.id, session_id=session_id
|
|
)
|
|
assert path.read_bytes() == _png_bytes()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_survives_material_copy_failure(
|
|
storage: SessionStorage, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
media_root = tmp_path / "media"
|
|
manager = SessionManager(storage, inject_time_prefix=False, media_root=media_root)
|
|
parent = await manager.create("agent:main:main")
|
|
await storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=parent.session_id,
|
|
session_key=parent.session_key,
|
|
role="assistant",
|
|
content="hi",
|
|
token_count=1,
|
|
)
|
|
)
|
|
|
|
def _boom(*args: object, **kwargs: object) -> int:
|
|
raise RuntimeError("disk exploded mid-copy")
|
|
|
|
monkeypatch.setattr(ArtifactStore, "copy_session_artifacts", _boom)
|
|
|
|
# The copy raising must NOT abort the fork: the child is still created/committed.
|
|
child = await manager.branch(
|
|
"agent:main:main", "agent:main:direct:resilient", fork_transcript=True
|
|
)
|
|
assert child.forked_from_parent is True
|
|
assert await manager.get_session("agent:main:direct:resilient") is not None
|