diff --git a/.github/workflows/image-memory-validation.yml b/.github/workflows/image-memory-validation.yml deleted file mode 100644 index a1ff61278..000000000 --- a/.github/workflows/image-memory-validation.yml +++ /dev/null @@ -1,215 +0,0 @@ -name: Image Memory Lifecycle Validation - -on: - push: - branches: - - feat/optimize-image-memory-lifecycle - workflow_dispatch: - -jobs: - validate: - name: Image and memory validation (${{ matrix.os }}) - runs-on: ${{ matrix.os }} - timeout-minutes: 90 - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-latest, macos-latest] - - defaults: - run: - shell: bash - - steps: - - name: Checkout - uses: actions/checkout@v7 - - - name: Set up Python - uses: actions/setup-python@v7 - with: - python-version: '3.12' - - - name: Install uv - run: | - python -m pip install --upgrade pip - python -m pip install uv - - - name: Install dependencies - run: uv sync - - - name: Run repository ruff checks - run: | - uv run ruff format --check . - uv run ruff check . - - - name: Run focused tests - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results" - uv run pytest -q \ - tests/unit/test_image_preparation_budget.py \ - tests/unit/test_image_media_store.py \ - tests/test_media_utils.py \ - tests/test_computer_fs_tools.py \ - tests/test_tool_loop_agent_runner.py \ - tests/agent/test_context_manager.py \ - tests/unit/test_astr_main_agent.py \ - tests/unit/test_context_image_budget.py \ - tests/unit/test_image_source_preparation.py \ - tests/unit/test_image_memory_benchmark.py \ - tests/unit/test_image_lifecycle_workloads.py \ - tests/unit/test_image_screenshot_encoding.py \ - tests/unit/test_image_source_integration.py \ - tests/unit/test_provider_image_references.py \ - tests/test_agent_runner_media_resolver.py \ - tests/unit/test_provider_request_retry.py \ - tests/unit/test_manage_image_history.py \ - tests/unit/test_conversation_media_api.py \ - tests/test_openai_source.py \ - tests/test_openai_responses_source.py \ - tests/test_gemini_source.py \ - tests/test_anthropic_source.py \ - tests/test_anthropic_kimi_code_provider.py \ - tests/test_deerflow_agent_runner.py \ - tests/unit/test_cron_context_compression.py \ - tests/unit/test_group_chat_context_wiring.py \ - tests/test_platform_image_format_preservation.py \ - tests/test_webchat_queue_lifecycle.py \ - tests/test_webchat_upload_image_format.py \ - --junitxml="$RUNNER_TEMP/image-memory-validation/results/tests.xml" - - - name: Generate reproducible fixtures outside measured processes - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/fixtures" - uv run python scripts/image_memory_bench/generate_fixtures.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures" \ - --seeds 7 19 43 - - - name: Record versions and fixture manifest - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results" - uv run python scripts/image_memory_bench/baseline_manifest.py \ - "$RUNNER_TEMP/image-memory-validation/results/manifest.json" \ - --fixtures "$RUNNER_TEMP/image-memory-validation/fixtures" - - - name: Run single-image memory workloads for every paired seed - run: | - set +e - result="$RUNNER_TEMP/image-memory-validation/results/measurements.jsonl" - failures=0 - for seed in 7 19 43; do - while IFS= read -r image; do - uv run python scripts/image_memory_bench/measure.py \ - "$image" "$result" \ - --repo "$GITHUB_WORKSPACE" \ - --csv "$RUNNER_TEMP/image-memory-validation/results/measurements.csv" \ - --workload "$seed-$(basename "$image")" - code=$? - if [ "$code" -ne 0 ]; then - failures=$((failures + 1)) - fi - done < <(find "$RUNNER_TEMP/image-memory-validation/fixtures/$seed" \ - -type f \( -name '*ordinary*' -o -name '*stress*' \) | sort) - done - printf '{"measurement_failures":%s}\n' "$failures" \ - > "$RUNNER_TEMP/image-memory-validation/results/measurement-summary.json" - - - name: Run Python allocation diagnostic workload - run: | - uv run python scripts/image_memory_bench/measure.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/7/png-stress.png" \ - "$RUNNER_TEMP/image-memory-validation/results/trace-python.jsonl" \ - --repo "$GITHUB_WORKSPACE" \ - --trace-python \ - --csv "$RUNNER_TEMP/image-memory-validation/results/trace-python.csv" \ - --workload trace-python-png-stress - - - name: Run paired history loading workloads - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results/history" - for seed in 7 19 43; do - uv run python scripts/image_memory_bench/history_benchmark.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/$seed/manifest.json" \ - "$RUNNER_TEMP/image-memory-validation/results/history/$seed.jsonl" \ - --repo "$GITHUB_WORKSPACE" \ - --images 50 \ - --keep-turns 25 \ - --repeats 10 \ - --cold-repeats 10 \ - --warm-repeats 10 \ - --distinct \ - --csv "$RUNNER_TEMP/image-memory-validation/results/history/$seed.csv" \ - --rss-jsonl "$RUNNER_TEMP/image-memory-validation/results/history/$seed.rss.jsonl" - done - - - name: Run fixed-window concurrent lifecycle and restart workload - run: | - lifecycle_root="$RUNNER_TEMP/image-memory-validation/lifecycle" - mkdir -p "$lifecycle_root" - uv run python scripts/image_memory_bench/prepare_lifecycle_fixture.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/7/manifest.json" \ - "$lifecycle_root/input" \ - --session-count 4 \ - --history-turns 50 - lifecycle_manifest="$lifecycle_root/input/lifecycle-manifest.json" - database=$(uv run python -c \ - 'import json,sys; print(json.load(open(sys.argv[1]))["database"])' \ - "$lifecycle_manifest") - media_root=$(uv run python -c \ - 'import json,sys; print(json.load(open(sys.argv[1]))["media_root"])' \ - "$lifecycle_manifest") - conversation_ids=$(uv run python -c \ - 'import json,sys; print(" ".join(json.load(open(sys.argv[1]))["conversation_ids"]))' \ - "$lifecycle_manifest") - command=( - uv run python scripts/image_memory_bench/lifecycle_workloads.py - "$lifecycle_root/result.json" - --repo "$GITHUB_WORKSPACE" - --database "$database" - --media-root "$media_root" - --session-count 4 - --request-count 200 - --window-turns 25 - --rss-jsonl "$lifecycle_root/result.rss.jsonl" - ) - for conversation_id in $conversation_ids; do - command+=(--conversation-id "$conversation_id") - done - "${command[@]}" - - - name: Record ablation contract status without fabricating factors - run: | - uv run python scripts/image_memory_bench/ablation_runner.py \ - "$RUNNER_TEMP/image-memory-validation/results/ablation-plan.json" \ - --baseline-repo "$GITHUB_WORKSPACE" \ - --candidate-repo "$GITHUB_WORKSPACE" - - - name: Upload raw validation artifacts - if: always() - uses: actions/upload-artifact@v4 - with: - name: image-memory-validation-${{ matrix.os }} - path: ${{ runner.temp }}/image-memory-validation - if-no-files-found: error - - dashboard: - name: Dashboard typecheck - runs-on: ubuntu-latest - timeout-minutes: 20 - defaults: - run: - working-directory: dashboard - steps: - - name: Checkout - uses: actions/checkout@v7 - - name: Set up Node - uses: actions/setup-node@v7 - with: - node-version: 22 - cache: pnpm - cache-dependency-path: dashboard/pnpm-lock.yaml - - name: Enable pnpm - run: corepack enable - - name: Install dashboard dependencies - run: pnpm install --frozen-lockfile - - name: Run dashboard typecheck - run: pnpm typecheck diff --git a/scripts/image_memory_bench/ablation_runner.py b/scripts/image_memory_bench/ablation_runner.py deleted file mode 100644 index 9c8b08abb..000000000 --- a/scripts/image_memory_bench/ablation_runner.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Run the real 2^3 image experiment only when production factor flips exist. - -The current checkout does not expose independent A, B, and C experiment -switches. This runner therefore refuses to fabricate a matrix and emits an -auditable blocked manifest until each factor is supplied by a real checkout. -""" - -from __future__ import annotations - -import argparse -import itertools -import json -import subprocess -from pathlib import Path - - -def git_identity(repo: Path) -> dict[str, str]: - """Return commit and dirty diff identity for a checkout.""" - commit = subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=repo, text=True - ).strip() - diff = subprocess.check_output(["git", "diff", "--binary"], cwd=repo) - import hashlib - - return { - "commit": commit, - "working_tree_diff_sha256": hashlib.sha256(diff).hexdigest(), - } - - -def main() -> int: - """Validate factor provenance and write a blocked or runnable plan.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--baseline-repo", type=Path, required=True) - parser.add_argument("--candidate-repo", type=Path, required=True) - parser.add_argument("--factor-entrypoint", type=Path) - args = parser.parse_args() - factors = { - "A": "unavailable: no independent production source-preparation entrypoint supplied", - "B": "unavailable: no independent production storage/lazy-materialization entrypoint supplied", - "C": "unavailable: no independent production compression-algorithm entrypoint supplied", - } - entrypoint_ok = ( - args.factor_entrypoint is not None and args.factor_entrypoint.is_file() - ) - plan = { - "status": "blocked" if not entrypoint_ok else "requires_factor_contract", - "experiment": "real-2^3-factorial", - "factors": factors, - "combinations": ["".join(bits) for bits in itertools.product("0A", repeat=0)], - "baseline": git_identity(args.baseline_repo), - "candidate": git_identity(args.candidate_repo), - "warmup_scope": "sdk_send_only; not warm history", - "reason": "Do not execute or label combinations until A/B/C independently select real production paths.", - } - plan["combinations"] = [ - "".join(bits) for bits in itertools.product("0A", "0B", "0C") - ] - args.output.write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") - print(json.dumps({"status": plan["status"], "output": str(args.output)})) - return 0 if plan["status"] == "blocked" else 2 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/baseline_manifest.py b/scripts/image_memory_bench/baseline_manifest.py deleted file mode 100644 index 1710c775c..000000000 --- a/scripts/image_memory_bench/baseline_manifest.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Record the exact code and dependency baseline for image experiments.""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.metadata -import json -import platform -import subprocess -import sys -from pathlib import Path - - -def main() -> int: - """Write a reproducibility manifest without reading application data.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--fixtures", type=Path) - args = parser.parse_args() - commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() - packages = sorted( - f"{dist.metadata['Name']}=={dist.version}" - for dist in importlib.metadata.distributions() - if dist.metadata.get("Name") - ) - manifest = { - "commit": commit, - "working_tree_diff_sha256": hashlib.sha256( - subprocess.check_output(["git", "diff", "--binary"], cwd=Path.cwd()) - ).hexdigest(), - "python": sys.version, - "platform": platform.platform(), - "packages": packages, - } - if args.fixtures: - files = [] - for path in sorted(args.fixtures.rglob("*")): - if path.is_file() and path.name != "manifest.json": - files.append( - { - "path": str(path.relative_to(args.fixtures)), - "bytes": path.stat().st_size, - "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), - } - ) - manifest["fixtures"] = files - args.output.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") - print(json.dumps({"commit": commit, "package_count": len(packages)})) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/generate_fixtures.py b/scripts/image_memory_bench/generate_fixtures.py deleted file mode 100644 index 74c8cc9a3..000000000 --- a/scripts/image_memory_bench/generate_fixtures.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Generate seeded fixtures outside the measured application process.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import random -from pathlib import Path - -from PIL import Image, ImageDraw - - -def generate(output: Path, seed: int) -> list[dict]: - """Create ten inputs for one paired experiment seed. - - Args: - output: New directory for this seed's fixtures. - seed: Stable seed for pixel generation. - - Returns: - On-disk hashes and decoded image metadata. - """ - output.mkdir(parents=True, exist_ok=False) - rng = random.Random(seed) - records = [] - for kind in ("png", "jpeg", "webp", "gif", "screenshot"): - for stress in (False, True): - name = f"{kind}-{'stress' if stress else 'ordinary'}" - options = {} - if kind in ("png", "webp"): - size = (1280, 1280) if stress else (320, 240) - image = ( - Image.frombytes("RGBA", size, rng.randbytes(size[0] * size[1] * 4)) - if stress - else Image.new("RGBA", size, (40, 120, 200, 127)) - ) - image_format = kind.upper() - options = {"lossless": stress} if kind == "webp" else {} - elif kind == "jpeg": - size = (4000, 3000) if stress else (640, 480) - image = Image.frombytes( - "RGB", size, rng.randbytes(size[0] * size[1] * 3) - ) - image_format = "JPEG" - options = {"quality": 95 if stress else 75} - if stress: - exif = image.getexif() - exif[274] = 6 - options["exif"] = exif - elif kind == "gif": - image = Image.new("RGB", (320, 240), (seed % 255, 0, 255)) - image_format = "GIF" - frames = [] - for frame in range(24 if stress else 3): - extra = Image.new("RGB", image.size, (frame * 9, 80, 20)) - ImageDraw.Draw(extra).rectangle( - (frame * 5, 10, frame * 5 + 30, 50), fill="white" - ) - frames.append(extra) - options = { - "save_all": True, - "append_images": frames, - "duration": 80, - "loop": 0, - } - else: - image = Image.new( - "RGB", (3840, 2160) if stress else (1920, 1080), "#202124" - ) - image_format = "PNG" - draw = ImageDraw.Draw(image) - for row in range(image.height // 20): - draw.text( - (100, row * 20), - f"coordinate (100, {row * 20}) seed={seed} command --verbose " - * 4, - fill="#e8eaed", - ) - draw.rectangle((20, 20, 70, 70), outline="red", width=3) - path = output / f"{name}.{image_format.lower()}" - try: - image.save(path, image_format, **options) - finally: - image.close() - for frame_image in options.get("append_images", []): - frame_image.close() - # Reopen the saved bytes: in-memory Image.format and frame counts - # do not describe what the encoder actually persisted. - with Image.open(path) as saved: - record = { - "path": path.name, - "kind": kind, - "stress": stress, - "seed": seed, - "format": saved.format, - "width": saved.width, - "height": saved.height, - "mode": saved.mode, - "frames": getattr(saved, "n_frames", 1), - "orientation": saved.getexif().get(274, 1), - "has_alpha": saved.mode in ("RGBA", "LA") - or "transparency" in saved.info, - "bytes": path.stat().st_size, - "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), - } - records.append(record) - (output / "manifest.json").write_text( - json.dumps(records, indent=2) + "\n", encoding="utf-8" - ) - return records - - -def main() -> int: - """Generate three paired seeds unless explicitly overridden.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--seeds", nargs="+", type=int, default=[7, 19, 43]) - args = parser.parse_args() - for seed in args.seeds: - records = generate(args.output / str(seed), seed) - print(json.dumps({"seed": seed, "count": len(records)})) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/history_benchmark.py b/scripts/image_memory_bench/history_benchmark.py deleted file mode 100644 index 0e9563bf6..000000000 --- a/scripts/image_memory_bench/history_benchmark.py +++ /dev/null @@ -1,461 +0,0 @@ -"""Measure SQLite history loading and real SDK request bodies for B-only.""" - -# The compact benchmark runner keeps subprocess orchestration visibly linear. -# ruff: noqa: E701, E702 - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import csv -import hashlib -import json -import mimetypes -import os -import subprocess -import sys -import tempfile -import threading -import time -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import psutil - - -async def prepare_database( - db_path: Path, - media_root: Path, - fixture: Path, - images: int, - distinct: bool, - mode: str, -) -> None: - """Create one SQLite input outside the measured child.""" - sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - persist_inline_image_refs, - ) - - rows = json.loads(fixture.read_text()) - paths = [fixture.parent / row["path"] for row in rows] - history = [] - raw = [] - for index in range(images): - path = paths[index % len(paths)] if distinct else paths[0] - # Distinct workloads use separately generated, valid image files listed - # by the fixture manifest. Never mutate encoded bytes into invalid images. - data = path.read_bytes() - mime = mimetypes.guess_type(path.name)[0] or "image/png" - raw.append( - { - "type": "image_url", - "image_url": { - "url": f"data:{mime};base64,{base64.b64encode(data).decode()}", - "detail": "high", - }, - } - ) - for index, part in enumerate(raw): - history.extend( - [ - {"role": "user", "content": [part]}, - {"role": "assistant", "content": f"ack {index}"}, - ] - ) - history.extend( - [ - {"role": "user", "content": "final text"}, - {"role": "assistant", "content": "final response"}, - ] - ) - if mode == "reference": - history = persist_inline_image_refs(history, ImageMediaStore(media_root)) - db = SQLiteDatabase(str(db_path)) - await db.initialize() - record = await db.create_conversation( - user_id="bench", platform_id="bench", content=history, title="benchmark" - ) - (db_path.parent / "conversation-id").write_text(record.conversation_id) - await db.engine.dispose() - - -def count_images(value: object) -> int: - """Count data image URLs in a decoded SDK body.""" - if isinstance(value, dict): - return sum(count_images(item) for item in value.values()) - if isinstance(value, list): - return sum(count_images(item) for item in value) - return int(isinstance(value, str) and value.startswith("data:image/")) - - -async def run_child(args: argparse.Namespace) -> dict: - """Read through SQLite and ConversationManager before selecting context.""" - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.agent.context.config import ContextConfig - from astrbot.core.agent.context.manager import ContextManager - from astrbot.core.agent.message import Message - from astrbot.core.conversation_mgr import ConversationManager - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - ) - - phases = {} - - def mark(phase: str) -> None: - with args.marker.open("a", encoding="utf-8") as marker: - marker.write( - json.dumps({"phase": phase, "monotonic": time.monotonic()}) + "\n" - ) - - mark("db_and_convert") - started = time.perf_counter() - db = SQLiteDatabase(str(args.database)) - await db.initialize() - conversation = await ConversationManager(db).get_conversation( - "bench", args.conversation_id - ) - phases["db_and_convert_ms"] = (time.perf_counter() - started) * 1000 - mark("json_load_bind") - started = time.perf_counter() - history = json.loads(conversation.history) - messages = [Message.model_validate(item) for item in history] - phases["json_load_bind_ms"] = (time.perf_counter() - started) * 1000 - mark("selection") - started = time.perf_counter() - selected = await ContextManager( - ContextConfig(enforce_max_turns=args.keep_turns) - ).process(messages) - phases["selection_ms"] = (time.perf_counter() - started) * 1000 - mark("materialize") - selected = await materialize_image_media_refs( - selected, ImageMediaStore(args.media_root) - ) - payload = [message.model_dump() for message in selected] - mark("sdk") - async with AsyncOpenAI( - api_key="local", base_url=args.endpoint, max_retries=0 - ) as client: - for _ in range(args.warmup): - await client.chat.completions.create(model="local", messages=payload) - await client.chat.completions.create(model="local", messages=payload) - await db.engine.dispose() - return { - "phases_ms": phases, - "selected_messages": len(payload), - "selected_images": count_images(payload), - } - - -def main() -> int: - """Prepare paired databases and monitor fresh children.""" - parser = argparse.ArgumentParser() - parser.add_argument("fixture", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--images", type=int, default=50) - parser.add_argument("--keep-turns", type=int, default=25) - parser.add_argument("--repeats", type=int, default=10) - parser.add_argument("--cold-repeats", type=int) - parser.add_argument("--warm-repeats", type=int) - parser.add_argument("--warmup", type=int, default=3) - parser.add_argument("--distinct", action="store_true") - parser.add_argument("--csv", type=Path) - parser.add_argument("--rss-jsonl", type=Path) - parser.add_argument("--child", action="store_true") - parser.add_argument("--database", type=Path) - parser.add_argument("--media-root", type=Path) - parser.add_argument("--history-mode", choices=("inline", "reference")) - parser.add_argument("--conversation-id", default="") - parser.add_argument("--endpoint") - parser.add_argument("--body-file", type=Path) - parser.add_argument("--marker", type=Path) - args = parser.parse_args() - cold_repeats = args.cold_repeats if args.cold_repeats is not None else args.repeats - warm_repeats = args.warm_repeats if args.warm_repeats is not None else args.repeats - if min(args.repeats, cold_repeats, warm_repeats, args.images, args.keep_turns) < 1: - parser.error("repetition, image, and context-window values must be positive") - fixture_rows = json.loads(args.fixture.read_text()) - max_fixture_bytes = max( - ( - row.get("bytes") or row.get("prepared_bytes") or row.get("source_bytes", 0) - for row in fixture_rows - ), - default=0, - ) - budget_class = ( - "over-4MiB-source-stress" - if 4 * ((max_fixture_bytes + 2) // 3) > 4 * 1024 * 1024 - else "representative-within-4MiB-source" - ) - if args.child: - if ( - args.database is None - or args.media_root is None - or args.endpoint is None - or args.body_file is None - or args.marker is None - or args.history_mode is None - ): - parser.error( - "child requires database, media root, endpoint, body file, marker, and history mode" - ) - result = asyncio.run(run_child(args)) - if args.body_file.exists(): - metadata = json.loads(args.body_file.read_text()) - result.update( - { - "http_json_bytes": metadata["bytes"], - "http_body_sha256": metadata["sha256"], - "http_images": metadata["images"], - } - ) - else: - result.update( - { - "http_json_bytes": None, - "http_body_sha256": None, - "http_images": 0, - } - ) - if not result["selected_images"] or not result["http_images"]: - raise RuntimeError("actual SDK body contains no images") - args.output.write_text(json.dumps(result) + "\n") - return 0 - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text("", encoding="utf-8") - csv_file = args.csv.open("w", newline="", encoding="utf-8") if args.csv else None - writer = None - received_body = {"path": None} - rss_path = args.output.with_name(args.output.stem + ".rss.jsonl") - if args.rss_jsonl: - rss_path = args.rss_jsonl - rss_path.parent.mkdir(parents=True, exist_ok=True) - rss_file = rss_path.open("w", encoding="utf-8") - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - data = self.rfile.read(int(self.headers.get("Content-Length", "0"))) - body = json.loads(data) - Path(received_body["path"]).write_text( - json.dumps( - { - "bytes": len(data), - "sha256": hashlib.sha256(data).hexdigest(), - "images": count_images(body), - } - ) - ) - response = b'{"choices":[{"message":{"content":"ok"}}]}' - self.send_response(200) - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, *_args): - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target=server.serve_forever, daemon=True).start() - endpoint = f"http://127.0.0.1:{server.server_port}/v1" - try: - for mode in ("inline", "reference"): - for repeat in range(cold_repeats + warm_repeats): - with tempfile.TemporaryDirectory( - prefix="astrbot-history-input-" - ) as temp: - root = Path(temp) - db_path = root / f"{mode}.db" - media = root / mode - asyncio.run( - prepare_database( - db_path, - media, - args.fixture, - args.images, - args.distinct, - mode, - ) - ) - conversation_id = (db_path.parent / "conversation-id").read_text() - child_output = root / "child.jsonl" - body_file = root / "http-metadata.json" - received_body["path"] = body_file - warmup = 0 if repeat < cold_repeats else args.warmup - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(args.fixture), - str(child_output), - "--child", - "--repo", - str(args.repo), - "--database", - str(db_path), - "--media-root", - str(media), - "--conversation-id", - conversation_id, - "--keep-turns", - str(args.keep_turns), - "--history-mode", - mode, - "--warmup", - str(warmup), - "--endpoint", - endpoint, - "--body-file", - str(body_file), - "--marker", - str(root / "marker.jsonl"), - ] - proc = subprocess.Popen( - command, - cwd=args.repo, - env={**os.environ, "ASTRBOT_ROOT": str(root)}, - ) - monitor = psutil.Process(proc.pid) - samples = [] - started = time.monotonic() - limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) - stop_reason = None - while proc.poll() is None: - try: - now = time.monotonic() - rss = monitor.memory_info().rss - private_bytes = None - try: - private_bytes = monitor.memory_full_info().uss - except (AttributeError, psutil.AccessDenied): - pass - samples.append( - { - "seconds": now - started, - "monotonic": now, - "rss": rss, - "private_bytes": private_bytes, - } - ) - if now - started > 120: - stop_reason = "timeout" - elif rss > limit: - stop_reason = "memory_limit" - if stop_reason: - proc.kill() - break - except psutil.NoSuchProcess: - break - time.sleep(0.01) - proc.wait() - if child_output.exists(): - result = json.loads(child_output.read_text()) - else: - result = { - "status": "error", - "error": stop_reason or "child_failed", - } - marker_file = root / "marker.jsonl" - markers = ( - [ - json.loads(line) - for line in marker_file.read_text().splitlines() - ] - if marker_file.exists() - else [] - ) - for sample in samples: - eligible = [ - marker - for marker in markers - if marker["monotonic"] <= sample["monotonic"] - ] - sample["phase"] = ( - max(eligible, key=lambda marker: marker["monotonic"])[ - "phase" - ] - if eligible - else "pre-start" - ) - result.update( - { - "mode": mode, - "repeat": repeat, - "cold": repeat < cold_repeats, - "warmup_mode": "sdk_warmup" if warmup else "cold", - "warmup_scope": "sdk_send_only", - "workload": f"sqlite-history-{args.images}-{'distinct' if args.distinct else 'same'}", - "budget_class": budget_class, - "rss_high_water_bytes": max( - (sample["rss"] for sample in samples), default=None - ), - "private_high_water_bytes": max( - ( - sample["private_bytes"] - for sample in samples - if sample["private_bytes"] is not None - ), - default=None, - ), - "stop_reason": stop_reason, - "rss_sample_count": len(samples), - "rss_jsonl": str(rss_path), - "returncode": proc.returncode, - "status": ( - "ok" - if proc.returncode == 0 and stop_reason is None - else "error" - ), - } - ) - for sample in samples: - rss_file.write( - json.dumps( - { - "mode": mode, - "repeat": repeat, - **sample, - } - ) - + "\n" - ) - rss_file.flush() - with args.output.open("a", encoding="utf-8") as stream: - stream.write(json.dumps(result) + "\n") - if csv_file: - fields = [ - "mode", - "repeat", - "cold", - "workload", - "rss_high_water_bytes", - "private_high_water_bytes", - "rss_sample_count", - "status", - "returncode", - "http_json_bytes", - "http_body_sha256", - "http_images", - ] - if writer is None: - writer = csv.DictWriter(csv_file, fieldnames=fields) - writer.writeheader() - writer.writerow({field: result.get(field) for field in fields}) - finally: - server.shutdown() - rss_file.close() - if csv_file: - csv_file.close() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/lifecycle_workloads.py b/scripts/image_memory_bench/lifecycle_workloads.py deleted file mode 100644 index 75c80846e..000000000 --- a/scripts/image_memory_bench/lifecycle_workloads.py +++ /dev/null @@ -1,371 +0,0 @@ -"""Run real history lifecycle workloads in an isolated, measured child process.""" - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import hashlib -import json -import subprocess -import sys -import tempfile -import threading -import time -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Any - -import psutil - -REQUESTS = 200 -SESSIONS = 4 -WINDOW_TURNS = 25 - - -def _count_images(value: Any) -> int: - if isinstance(value, dict): - return sum(_count_images(item) for item in value.values()) - if isinstance(value, list): - return sum(_count_images(item) for item in value) - return int(isinstance(value, str) and value.startswith("data:image/")) - - -def _image_hashes(value: Any) -> list[str]: - if isinstance(value, dict): - if value.get("type") == "image_url" and isinstance( - value.get("image_url"), dict - ): - encoded = value["image_url"].get("url", "").split(",", 1)[-1] - try: - return [hashlib.sha256(base64.b64decode(encoded)).hexdigest()] - except Exception: - return [] - return [item for child in value.values() for item in _image_hashes(child)] - if isinstance(value, list): - return [item for child in value for item in _image_hashes(child)] - return [] - - -async def _child(args: argparse.Namespace) -> dict[str, Any]: - """Execute the production database, context, materialization, and SDK path.""" - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.agent.context.config import ContextConfig - from astrbot.core.agent.context.manager import ContextManager - from astrbot.core.agent.message import Message - from astrbot.core.conversation_mgr import ConversationManager - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - ) - - async def load_payload( - database: SQLiteDatabase, conversation_id: str - ) -> tuple[list, str, list]: - conversation = await ConversationManager(database).get_conversation( - "bench", conversation_id - ) - messages = [ - Message.model_validate(item) for item in json.loads(conversation.history) - ] - selected = await ContextManager( - ContextConfig(enforce_max_turns=args.window_turns) - ).process(messages) - selected = await materialize_image_media_refs( - selected, ImageMediaStore(args.media_root) - ) - payload = [message.model_dump() for message in selected] - semantic_hash = hashlib.sha256( - json.dumps(payload, sort_keys=True, ensure_ascii=False).encode() - ).hexdigest() - return payload, semantic_hash, messages - - async def session( - client: AsyncOpenAI, session_id: str, request_count: int - ) -> dict[str, Any]: - database = SQLiteDatabase(str(args.database)) - await database.initialize() - semantic_hashes: list[str] = [] - image_counts: list[int] = [] - image_hash_orders: list[list[str]] = [] - for turn in range(request_count): - _, _, messages = await load_payload(database, session_id) - image_refs = [ - part.model_dump() if hasattr(part, "model_dump") else part - for message in messages - for part in ( - message.content if isinstance(message.content, list) else [] - ) - if (part.model_dump() if hasattr(part, "model_dump") else part).get( - "type" - ) - == "image_media_ref" - ] - if not image_refs: - raise RuntimeError( - "lifecycle fixture has no persisted image references" - ) - messages.extend( - [ - Message(role="user", content=[image_refs[turn % len(image_refs)]]), - Message(role="assistant", content=f"ack {turn}"), - ] - ) - await ConversationManager(database).update_conversation( - "bench", - session_id, - history=[message.model_dump() for message in messages], - ) - payload, semantic_hash, _ = await load_payload(database, session_id) - semantic_hashes.append(semantic_hash) - image_counts.append(_count_images(payload)) - image_hash_orders.append(_image_hashes(payload)) - if not image_hash_orders[-1] or image_counts[-1] > args.window_turns: - raise RuntimeError("active image window is empty or unbounded") - await client.chat.completions.create(model="local", messages=payload) - del payload - await database.engine.dispose() - return { - "session": session_id, - "requests": request_count, - "image_counts": image_counts, - "image_hash_orders": image_hash_orders, - "semantic_hashes": semantic_hashes, - } - - if args.restart_check: - checks = [] - database = SQLiteDatabase(str(args.database)) - await database.initialize() - for session_id in args.conversation_ids: - payload, semantic_hash, _ = await load_payload(database, session_id) - checks.append( - { - "session": session_id, - "images": _count_images(payload), - "semantic_hash": semantic_hash, - "image_hash_order": _image_hashes(payload), - } - ) - await database.engine.dispose() - return {"restart_checks": checks} - - async with AsyncOpenAI( - api_key="local", base_url=args.endpoint, max_retries=0 - ) as client: - session_count = len(args.conversation_ids) - base_requests, remainder = divmod(args.request_count, session_count) - request_counts = [ - base_requests + (index < remainder) for index in range(session_count) - ] - results = await asyncio.gather( - *( - session(client, cid, request_count) - for cid, request_count in zip(args.conversation_ids, request_counts) - ) - ) - return { - "sessions": results, - "total_requests": sum(request_counts), - "request_count_per_session": request_counts, - "window_turns": args.window_turns, - } - - -def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--child", action="store_true") - parser.add_argument("--database", type=Path) - parser.add_argument("--media-root", type=Path) - parser.add_argument( - "--conversation-id", dest="conversation_ids", action="append", default=[] - ) - parser.add_argument("--endpoint") - parser.add_argument("--rss-jsonl", type=Path) - parser.add_argument("--restart-check", action="store_true") - parser.add_argument("--session-count", type=int, choices=(1, 4), default=4) - parser.add_argument("--request-count", type=int, default=REQUESTS) - parser.add_argument("--window-turns", type=int, default=WINDOW_TURNS) - args = parser.parse_args() - if args.request_count < 1 or args.window_turns < 1: - parser.error("request count and window turns must be positive") - if args.child: - if ( - args.database is None - or args.media_root is None - or args.endpoint is None - or len(args.conversation_ids) != args.session_count - ): - parser.error( - "child requires database, media root, endpoint, and the requested conversation IDs" - ) - result = asyncio.run(_child(args)) - args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") - return 0 - if len(args.conversation_ids) != args.session_count: - parser.error( - f"exactly {args.session_count} --conversation-id values are required" - ) - if args.database is None or args.media_root is None: - parser.error("--database and --media-root are required") - - requests: list[dict[str, Any]] = [] - - class Sink(BaseHTTPRequestHandler): - def do_POST(self) -> None: - size = int(self.headers.get("Content-Length", "0")) - body = self.rfile.read(size) - requests.append( - { - "bytes": len(body), - "sha256": hashlib.sha256(body).hexdigest(), - "images": _count_images(json.loads(body)), - } - ) - response = b'{"choices":[{"message":{"content":"ok"}}]}' - self.send_response(200) - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, *_args: Any) -> None: - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Sink) - threading.Thread(target=server.serve_forever, daemon=True).start() - with tempfile.TemporaryDirectory(prefix="astrbot-lifecycle-") as run_dir: - child_output = Path(run_dir) / "child.json" - rss_path = args.rss_jsonl or args.output.with_name( - args.output.stem + ".rss.jsonl" - ) - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(child_output), - "--child", - "--repo", - str(args.repo.resolve()), - "--database", - str(args.database.resolve()), - "--media-root", - str(args.media_root.resolve()), - "--endpoint", - f"http://127.0.0.1:{server.server_port}/v1", - "--session-count", - str(args.session_count), - "--request-count", - str(args.request_count), - "--window-turns", - str(args.window_turns), - ] - for conversation_id in args.conversation_ids: - command.extend(["--conversation-id", conversation_id]) - started = time.monotonic() - process = subprocess.Popen(command, cwd=args.repo) - child = psutil.Process(process.pid) - stop_reason = None - with rss_path.open("w", encoding="utf-8") as rss: - while process.poll() is None: - current = child.memory_info().rss - private_bytes = None - try: - private_bytes = child.memory_full_info().uss - except (AttributeError, psutil.AccessDenied): - pass - rss.write( - json.dumps( - { - "seconds": time.monotonic() - started, - "rss": current, - "private_bytes": private_bytes, - } - ) - + "\n" - ) - rss.flush() - if ( - current > min(2 * 1024**3, psutil.virtual_memory().available // 4) - or time.monotonic() - started > 120 - ): - stop_reason = ( - "memory_limit" - if current - > min(2 * 1024**3, psutil.virtual_memory().available // 4) - else "timeout" - ) - process.kill() - break - time.sleep(0.01) - return_code = process.wait() - server.shutdown() - result = ( - json.loads(child_output.read_text()) - if child_output.exists() - else {"error": "child failed"} - ) - validation_errors = [] - restart_result = {"skipped": stop_reason or "child_failed"} - if stop_reason is None and return_code == 0: - restart_output = Path(run_dir) / "restart.json" - restart_command = command.copy() - restart_command[2] = str(restart_output) - restart_command.append("--restart-check") - restart_process = subprocess.run( - restart_command, cwd=args.repo, check=False - ) - restart_result = ( - json.loads(restart_output.read_text()) - if restart_output.exists() - else {"error": "restart child failed"} - ) - if restart_result.get("restart_checks"): - expected = { - item["session"]: item for item in result.get("sessions", []) - } - for check in restart_result["restart_checks"]: - prior = expected.get(check["session"]) - if ( - prior is None - or check["images"] != prior["image_counts"][-1] - or check["image_hash_order"] != prior["image_hash_orders"][-1] - ): - validation_errors.append( - "fresh-process restart changed the active image window" - ) - if restart_process.returncode != 0: - return_code = restart_process.returncode - result.update( - { - "returncode": return_code, - "stop_reason": stop_reason, - "request_count": len(requests), - "expected_request_count": args.request_count, - "requests": requests, - "rss_jsonl": str(rss_path), - "restart": restart_result, - } - ) - if any( - request["images"] <= 0 or request["images"] > args.window_turns - for request in requests - ): - validation_errors.append("sink observed an invalid active image window") - if len(requests) != args.request_count: - validation_errors.append("sink did not observe the expected request count") - if validation_errors: - result["validation_errors"] = validation_errors - return_code = return_code or 1 - result["returncode"] = return_code - args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") - return return_code - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/measure.py b/scripts/image_memory_bench/measure.py deleted file mode 100644 index 74c4bdb26..000000000 --- a/scripts/image_memory_bench/measure.py +++ /dev/null @@ -1,399 +0,0 @@ -"""Measure real image preparation and SDK requests in isolated processes.""" - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import csv -import hashlib -import json -import os -import subprocess -import sys -import tempfile -import threading -import time -import tracemalloc -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import psutil - - -def _base64_payload_metrics(encoded: str) -> tuple[int, str]: - """Hash an encoded payload in chunks without retaining decoded image bytes. - - Args: - encoded: Base64 payload without the data-URI prefix. - - Returns: - The decoded byte count and SHA-256 digest. - """ - digest = hashlib.sha256() - decoded_bytes = 0 - remainder = "" - chunk_size = 4 * 1024 * 1024 - for start in range(0, len(encoded), chunk_size): - chunk = remainder + encoded[start : start + chunk_size] - usable = len(chunk) - len(chunk) % 4 - if usable: - decoded = base64.b64decode(chunk[:usable]) - digest.update(decoded) - decoded_bytes += len(decoded) - remainder = chunk[usable:] - if remainder: - decoded = base64.b64decode(remainder) - digest.update(decoded) - decoded_bytes += len(decoded) - return decoded_bytes, digest.hexdigest() - - -async def run_child(args: argparse.Namespace) -> dict: - """Execute production preparation and request assembly in the chosen checkout. - - Args: - args: Explicit source, checkout, local server and diagnostic options. - - Returns: - Sizes, phase durations and resource counters without image content. - """ - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.provider.entities import ProviderRequest - - if args.trace_python: - tracemalloc.start() - process = psutil.Process() - baseline_rss = process.memory_info().rss - phases: dict[str, float] = {} - stage = "startup" - phase_marks = [(time.monotonic(), stage)] - failure: dict[str, str] | None = None - source_bytes = args.image.stat().st_size - source_digest = hashlib.sha256() - with args.image.open("rb") as source_stream: - for chunk in iter(lambda: source_stream.read(1024 * 1024), b""): - source_digest.update(chunk) - source_sha256 = source_digest.hexdigest() - prepared_bytes: int | None = None - prepared_sha256: str | None = None - started = time.perf_counter() - with tempfile.TemporaryDirectory(prefix="astrbot-image-runtime-") as runtime: - os.environ["ASTRBOT_ROOT"] = runtime - before_files = set(Path(runtime).rglob("*")) - try: - stage = "provider_prepare_and_assemble" - phase_marks.append((time.monotonic(), stage)) - tick = time.perf_counter() - request = ProviderRequest( - prompt="Describe this image.", - image_urls=[str(args.image.resolve())], - ) - context = await request.assemble_context() - phases["provider_prepare_and_assemble_ms"] = ( - time.perf_counter() - tick - ) * 1000 - image_part = next( - part - for part in context.get("content", []) - if part.get("type") == "image_url" - ) - data_url = image_part["image_url"]["url"] - stage = "benchmark_payload_verification" - phase_marks.append((time.monotonic(), stage)) - prepared_bytes, prepared_sha256 = _base64_payload_metrics( - data_url.split(",", 1)[1] - ) - stage = "sdk_send" - phase_marks.append((time.monotonic(), stage)) - async with AsyncOpenAI( - api_key="local-experiment-only", - base_url=args.endpoint, - max_retries=0, - ) as client: - for _ in range(args.warmup): - await client.chat.completions.create( - model="local", messages=[context] - ) - tick = time.perf_counter() - for _ in range(args.repeat): - await client.chat.completions.create( - model="local", messages=[context] - ) - phases["sdk_send_ms"] = (time.perf_counter() - tick) * 1000 - del context, request - except Exception as exc: - failure = {"stage": stage, "error_type": type(exc).__name__} - finally: - residual = [ - path - for path in Path(runtime).rglob("*") - if path.is_file() and path not in before_files - ] - result = { - "source_bytes": source_bytes, - "prepared_bytes": prepared_bytes, - "source_sha256": source_sha256, - "prepared_sha256": prepared_sha256, - "preserved_source_bytes": ( - prepared_sha256 == source_sha256 - if prepared_sha256 is not None - else None - ), - "base64_bytes": ( - 4 * ((prepared_bytes + 2) // 3) - if prepared_bytes is not None - else None - ), - "phase_ms": phases, - "phase_marks": phase_marks, - "last_stage": stage, - "failure": failure, - "baseline_rss": baseline_rss, - "post_request_rss": process.memory_info().rss, - "temporary_files_remaining": len(residual), - "temporary_bytes_remaining": sum( - path.stat().st_size for path in residual - ), - "elapsed_ms": (time.perf_counter() - started) * 1000, - } - if args.trace_python: - result["python_current_bytes"], result["python_peak_bytes"] = ( - tracemalloc.get_traced_memory() - ) - tracemalloc.stop() - if sys.platform != "win32": - import resource - - peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - result["os_peak_rss_bytes"] = ( - peak if sys.platform == "darwin" else peak * 1024 - ) - return result - - -def main() -> int: - """Run an isolated request and record measurements or a bounded failure.""" - parser = argparse.ArgumentParser() - parser.add_argument("image", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--trace-python", action="store_true") - parser.add_argument("--child", action="store_true") - parser.add_argument("--endpoint") - parser.add_argument("--repeat", type=int, default=1) - parser.add_argument("--warmup", type=int, default=0) - parser.add_argument("--csv") - parser.add_argument("--workload", default="single-image") - args = parser.parse_args() - if args.child: - try: - result = asyncio.run(run_child(args)) - result["status"] = "error" if result.get("failure") else "ok" - except Exception as error: - result = {"status": "error", "error_type": type(error).__name__} - args.output.write_text(json.dumps(result), encoding="utf-8") - return 0 if result["status"] == "ok" else 1 - - # The server measures actual SDK bytes without storing the request body. - request_sizes = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - length = int(self.headers.get("Content-Length", "0")) - remaining = length - while remaining: - chunk = self.rfile.read(min(65536, remaining)) - if not chunk: - break - remaining -= len(chunk) - request_sizes.append(length - remaining) - body = b'{"id":"local","object":"chat.completion","created":0,"model":"local","choices":[{"index":0,"message":{"role":"assistant","content":"fixture response"},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"completion_tokens":2,"total_tokens":102}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, *args): - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) - started = time.monotonic() - samples = [] - stop_reason = None - try: - with tempfile.TemporaryDirectory(prefix="astrbot-image-measure-") as run_dir: - child_result = Path(run_dir) / "child.json" - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(args.image.resolve()), - str(child_result), - "--child", - "--repo", - str(args.repo.resolve()), - "--endpoint", - f"http://127.0.0.1:{server.server_port}/v1", - ] - if args.repeat != 1: - command.extend(["--repeat", str(args.repeat)]) - if args.warmup: - command.extend(["--warmup", str(args.warmup)]) - command.extend(["--workload", args.workload]) - if args.trace_python: - command.append("--trace-python") - env = os.environ.copy() - env["ASTRBOT_ROOT"] = str(Path(run_dir) / "runtime") - with ( - (Path(run_dir) / "stdout").open("wb") as out, - (Path(run_dir) / "stderr").open("wb") as err, - ): - proc = subprocess.Popen( - command, cwd=args.repo, env=env, stdout=out, stderr=err - ) - process = psutil.Process(proc.pid) - try: - while proc.poll() is None: - try: - children = process.children(recursive=True) - rss = process.memory_info().rss - descendants_rss = sum( - p.memory_info().rss for p in children if p.is_running() - ) - info = process.memory_info() - samples.append( - { - "seconds": time.monotonic() - started, - "rss": rss, - "descendants_rss": descendants_rss, - "private_bytes": getattr(info, "private", None), - } - ) - if rss + descendants_rss > limit: - stop_reason = "memory_limit" - elif time.monotonic() - started > 120: - stop_reason = "timeout" - if stop_reason: - for child in children: - child.kill() - proc.kill() - break - except psutil.NoSuchProcess: - break - time.sleep(0.01) - finally: - if proc.poll() is None: - proc.kill() - proc.wait() - result = ( - json.loads(child_result.read_text()) - if child_result.exists() - else {"status": "error"} - ) - phase_peaks = {} - for sample in samples: - phase = "process_startup" - for timestamp, name in result.get("phase_marks", []): - if timestamp > started + sample["seconds"]: - break - phase = name - sample["phase"] = phase - phase_peaks[phase] = max(phase_peaks.get(phase, 0), sample["rss"]) - result.update( - { - "status": stop_reason or result["status"], - "workload": args.workload, - "phase_rss_high_water_bytes": phase_peaks, - "returncode": proc.returncode, - "rss_high_water_bytes": max( - (s["rss"] for s in samples), default=None - ), - "private_high_water_bytes": max( - ( - s["private_bytes"] - for s in samples - if s["private_bytes"] is not None - ), - default=None, - ), - "descendants_high_water_bytes": max( - (s["descendants_rss"] for s in samples), default=None - ), - "request_json_bytes": request_sizes, - "samples": samples, - "trace_python": args.trace_python, - "code_commit": subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=args.repo, text=True - ).strip(), - "working_tree_diff_sha256": hashlib.sha256( - subprocess.check_output( - ["git", "diff", "--binary"], cwd=args.repo - ) - ).hexdigest(), - "benchmark_source_sha256": hashlib.sha256( - Path(__file__).read_bytes() - ).hexdigest(), - "monitor_sample_count": len(samples), - } - ) - with args.output.open("a", encoding="utf-8") as output: - output.write(json.dumps(result) + "\n") - if args.csv: - fields = [ - "status", - "workload", - "repeat", - "warmup", - "rss_high_water_bytes", - "private_high_water_bytes", - "post_request_rss", - "temporary_files_remaining", - "temporary_bytes_remaining", - "request_json_bytes", - ] - csv_path = Path(args.csv) - write_header = not csv_path.exists() - with csv_path.open("a", newline="", encoding="utf-8") as csv_file: - writer = csv.DictWriter(csv_file, fieldnames=fields) - if write_header: - writer.writeheader() - writer.writerow( - { - "status": result.get("status"), - "workload": args.workload, - "repeat": args.repeat, - "warmup": args.warmup, - "rss_high_water_bytes": result.get("rss_high_water_bytes"), - "private_high_water_bytes": result.get( - "private_high_water_bytes" - ), - "post_request_rss": result.get("post_request_rss"), - "temporary_files_remaining": result.get( - "temporary_files_remaining" - ), - "temporary_bytes_remaining": result.get( - "temporary_bytes_remaining" - ), - "request_json_bytes": json.dumps( - result.get("request_json_bytes", []) - ), - } - ) - print(json.dumps({k: v for k, v in result.items() if k != "samples"})) - return 0 if result["status"] == "ok" else 1 - finally: - server.shutdown() - server.server_close() - server_thread.join() - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/migrate_history.py b/scripts/image_memory_bench/migrate_history.py deleted file mode 100644 index d2189d7c9..000000000 --- a/scripts/image_memory_bench/migrate_history.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Explicitly externalize or re-inline exported JSONL conversation histories. - -The input is never modified. Inspect the dry run before using --apply. Keep the -input alongside the media directory as the rollback backup. -""" - -from __future__ import annotations - -import argparse -import asyncio -import json -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - persist_inline_image_refs, -) - - -async def main() -> int: - """Transform an export into a new file without modifying live histories.""" - parser = argparse.ArgumentParser() - parser.add_argument("input", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument("--media-dir", type=Path, required=True) - parser.add_argument( - "--mode", choices=["externalize", "inline"], default="externalize" - ) - parser.add_argument("--apply", action="store_true") - parser.add_argument("--rollback", action="store_true") - args = parser.parse_args() - if args.input.resolve() == args.output.resolve() or args.output.exists(): - parser.error("Output must be a new file distinct from the rollback input") - if args.rollback: - if not args.apply: - parser.error("--rollback requires --apply") - # The input is the verified original export; rollback writes that exact - # export to a new path, so the original remains the recovery anchor. - args.mode = "inline" - store = ImageMediaStore(args.media_dir) - record_count = 0 - image_count = 0 - output = args.output.open("x", encoding="utf-8") if args.apply else None - try: - with args.input.open(encoding="utf-8") as source: - for line in source: - record = json.loads(line) - raw_history = record.get("history", []) - history = ( - json.loads(raw_history) - if isinstance(raw_history, str) - else raw_history - ) - if not isinstance(history, list): - raise ValueError("History must be a message list") - for message in history: - parts = message.get("content") - if isinstance(parts, list): - image_count += sum( - isinstance(part, dict) - and part.get("type") in {"image_url", "image_media_ref"} - for part in parts - ) - if args.apply: - if args.mode == "externalize" and not args.rollback: - history = persist_inline_image_refs(history, store) - else: - history = await materialize_image_media_refs( - history, store, strict=True - ) - record["history"] = ( - json.dumps(history, ensure_ascii=False) - if isinstance(raw_history, str) - else history - ) - output.write(json.dumps(record, ensure_ascii=False) + "\n") - record_count += 1 - except BaseException: - if output is not None: - output.close() - args.output.unlink(missing_ok=True) - raise - finally: - if output is not None: - output.close() - print( - json.dumps( - {"records": record_count, "images": image_count, "applied": args.apply} - ) - ) - return 0 - - -if __name__ == "__main__": - raise SystemExit(asyncio.run(main())) diff --git a/scripts/image_memory_bench/prepare_lifecycle_fixture.py b/scripts/image_memory_bench/prepare_lifecycle_fixture.py deleted file mode 100644 index b67f9452e..000000000 --- a/scripts/image_memory_bench/prepare_lifecycle_fixture.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Prepare a durable multi-session history fixture before measurement.""" - -from __future__ import annotations - -import argparse -import asyncio -import json -from copy import deepcopy -from pathlib import Path - - -async def prepare_fixture( - manifest_path: Path, - output: Path, - session_count: int, - history_turns: int, - include_stress: bool, -) -> dict: - """Create shared media and persisted conversations outside the measured child. - - Args: - manifest_path: Fixture manifest containing valid prepared image paths. - output: Fresh directory receiving the SQLite database and media objects. - session_count: Number of conversations to create. - history_turns: Number of image-bearing turns per conversation. - include_stress: Whether to include high-entropy stress fixtures. - - Returns: - A JSON-compatible manifest consumed by ``lifecycle_workloads.py``. - - Raises: - FileExistsError: The output already contains a database. - ValueError: The input manifest or requested sizes are invalid. - """ - if session_count < 1 or history_turns < 1: - raise ValueError("session count and history turns must be positive") - rows = json.loads(manifest_path.read_text(encoding="utf-8")) - if not isinstance(rows, list) or not rows: - raise ValueError("fixture manifest must contain at least one image") - if not include_stress: - rows = [row for row in rows if not row.get("stress")] - if not rows: - raise ValueError("fixture manifest must contain at least one image") - output.mkdir(parents=True, exist_ok=True) - database_path = output / "history.db" - if database_path.exists(): - raise FileExistsError(f"refusing to overwrite {database_path}") - - import sys - - sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ImageMediaStore - - media_root = output / "media" - store = ImageMediaStore(media_root) - refs = [] - for row in rows: - path = manifest_path.parent / row["path"] - refs.append(store.put(path.read_bytes(), detail="high")) - - history = [] - for turn in range(history_turns): - history.extend( - [ - { - "role": "user", - "content": [ - {"type": "text", "text": f"fixture turn {turn}"}, - refs[turn % len(refs)].model_dump(), - ], - }, - {"role": "assistant", "content": f"ack {turn}"}, - ] - ) - - database = SQLiteDatabase(str(database_path)) - await database.initialize() - conversation_ids = [] - for index in range(session_count): - conversation = await database.create_conversation( - user_id="bench", - platform_id="bench", - content=deepcopy(history), - title=f"lifecycle fixture {index}", - ) - conversation_ids.append(conversation.conversation_id) - await database.engine.dispose() - return { - "database": str(database_path.resolve()), - "media_root": str(media_root.resolve()), - "conversation_ids": conversation_ids, - "session_count": session_count, - "history_turns": history_turns, - "media_count": len(refs), - } - - -async def main_async(args: argparse.Namespace) -> int: - """Prepare the fixture and write its machine-readable manifest.""" - result = await prepare_fixture( - args.manifest, - args.output, - args.session_count, - args.history_turns, - args.include_stress, - ) - args.output.mkdir(parents=True, exist_ok=True) - output_manifest = args.output / "lifecycle-manifest.json" - output_manifest.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") - print(json.dumps(result)) - return 0 - - -def main() -> int: - """Parse fixture preparation arguments.""" - parser = argparse.ArgumentParser() - parser.add_argument("manifest", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument("--session-count", type=int, default=4) - parser.add_argument("--history-turns", type=int, default=50) - parser.add_argument( - "--include-stress", - action="store_true", - help="Include high-entropy stress fixtures in the long-running workload.", - ) - return asyncio.run(main_async(parser.parse_args())) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/manage_image_history.py b/scripts/manage_image_history.py deleted file mode 100644 index 4a3635bba..000000000 --- a/scripts/manage_image_history.py +++ /dev/null @@ -1,434 +0,0 @@ -#!/usr/bin/env python3 -"""Explicit offline migration, repair, rollback, and cleanup for image history.""" - -from __future__ import annotations - -import argparse -import base64 -import hashlib -import json -import shutil -import sqlite3 -import sys -import tempfile -from collections.abc import Iterator -from pathlib import Path -from typing import Any - -from astrbot.core.utils.image_media_store import ImageMediaRef, ImageMediaStore - - -def _digest(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _logical_digest(connection: sqlite3.Connection) -> str: - """Hash every table schema and row without loading the database in memory.""" - digest = hashlib.sha256() - tables = connection.execute( - "SELECT name, sql, type FROM sqlite_master WHERE type IN ('table','index','trigger','view') ORDER BY name" - ) - for name, sql, object_type in tables: - digest.update(json.dumps([name, sql], ensure_ascii=False).encode()) - if object_type != "table": - continue - rows = connection.execute( - f'SELECT * FROM "{name.replace(chr(34), chr(34) * 2)}"' - ) - for row in rows: - digest.update( - json.dumps( - list(row), ensure_ascii=False, default=str, separators=(",", ":") - ).encode() - ) - return digest.hexdigest() - - -def _manifest_media_name(name: str) -> tuple[str, str]: - path = Path(name) - if path.name != name or path.suffix not in {".bin", ".json"}: - raise ValueError("invalid media manifest filename") - media_id = path.stem - if len(media_id) != 64 or any(char not in "0123456789abcdef" for char in media_id): - raise ValueError("invalid media manifest media id") - return media_id, path.suffix - - -def _inline_parts(value: Any) -> Iterator[tuple[dict[str, Any], str, str]]: - if isinstance(value, list): - for item in value: - yield from _inline_parts(item) - elif isinstance(value, dict): - if value.get("type") == "image_url" and isinstance( - value.get("image_url"), dict - ): - url = value["image_url"].get("url") - if isinstance(url, str) and url.startswith("data:image/"): - header, encoded = url.split(",", 1) - yield value, header[5:].split(";", 1)[0], encoded - else: - for child in value.values(): - yield from _inline_parts(child) - - -def _replace_inline(value: Any, store: ImageMediaStore) -> tuple[Any, int, int]: - changed = 0 - bytes_migrated = 0 - - def visit(item: Any) -> Any: - nonlocal changed, bytes_migrated - if isinstance(item, list): - return [visit(child) for child in item] - if not isinstance(item, dict): - return item - if item.get("type") == "image_url" and isinstance(item.get("image_url"), dict): - image = item["image_url"] - url = image.get("url") - if isinstance(url, str) and url.startswith("data:image/"): - header, encoded = url.split(",", 1) - data = base64.b64decode(encoded, validate=True) - ref = store.put( - data, - header[5:].split(";", 1)[0], - image.get("detail"), - image.get("id"), - ) - changed += 1 - bytes_migrated += len(data) - return ref.model_dump() - return {key: visit(child) for key, child in item.items()} - - return visit(value), changed, bytes_migrated - - -def _connect(path: Path) -> sqlite3.Connection: - connection = sqlite3.connect(path) - connection.row_factory = sqlite3.Row - connection.execute("PRAGMA foreign_keys=ON") - return connection - - -def _backup(db: Path, media: Path, destination: Path) -> Path: - """Create a consistent database snapshot and a verified media snapshot.""" - destination.mkdir(parents=True, exist_ok=False) - snapshot = sqlite3.connect(destination / db.name) - readonly = sqlite3.connect(f"file:{db}?mode=ro", uri=True) - try: - readonly.backup(snapshot) - snapshot.commit() - finally: - readonly.close() - snapshot.close() - if media.exists(): - (destination / "media").mkdir() - manifest_media = [] - for path in media.iterdir(): - if path.is_symlink() or not path.is_file(): - raise RuntimeError("media backup refuses symlinks and non-files") - target = destination / "media" / path.name - shutil.copy2(path, target) - manifest_media.append({"name": path.name, "sha256": _digest(target)}) - else: - manifest_media = [] - manifest = { - "database": _digest(destination / db.name), - "media": manifest_media, - } - if (destination / "media").exists(): - manifest["media"] = sorted(manifest_media, key=lambda item: item["name"]) - (destination / "manifest.json").write_text( - json.dumps(manifest, sort_keys=True) + "\n" - ) - return destination - - -def _migrate(args: argparse.Namespace) -> dict[str, int]: - if args.apply and not args.offline: - raise SystemExit("--apply requires --offline: stop all history writes first") - db = Path(args.database).resolve() - media = Path(args.media).resolve() - connection = _connect(db) - lock = args.apply - try: - if lock: - connection.execute("BEGIN IMMEDIATE") - backup = _backup( - db, - media, - Path(args.backup) - if args.backup - else Path( - tempfile.mkdtemp(prefix="astrbot-image-history-backup-parent-") - ) - / "snapshot", - ) - query = ( - "SELECT inner_conversation_id, conversation_id, content FROM conversations" - ) - parameters: tuple[Any, ...] = () - if args.conversation_id: - query += ( - " WHERE conversation_id IN (" - + ",".join("?" for _ in args.conversation_id) - + ")" - ) - parameters = tuple(args.conversation_id) - total = changed = bytes_migrated = 0 - store = ImageMediaStore(media) - for row in connection.execute(query, parameters): - total += 1 - try: - content = ( - json.loads(row["content"]) - if isinstance(row["content"], str) - else row["content"] - ) - converted, count, amount = ( - _replace_inline(content, store) - if args.apply - else (content, sum(1 for _ in _inline_parts(content)), 0) - ) - except Exception as exc: # noqa: BLE001 - raise RuntimeError( - f"conversation {row['conversation_id']} cannot be parsed or verified" - ) from exc - changed += count - bytes_migrated += amount - if args.apply and count: - connection.execute( - "UPDATE conversations SET content=? WHERE inner_conversation_id=?", - ( - json.dumps(converted, ensure_ascii=False), - row["inner_conversation_id"], - ), - ) - if args.apply: - post_migration_digest = _logical_digest(connection) - connection.commit() - manifest_path = backup / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["post_migration_logical_digest"] = post_migration_digest - manifest_path.write_text(json.dumps(manifest, sort_keys=True) + "\n") - print( - json.dumps( - { - "backup": str(backup), - "records": total, - "images": changed, - "bytes": bytes_migrated, - } - ) - ) - else: - connection.rollback() - print(json.dumps({"dry_run": True, "records": total, "images": changed})) - return {"records": total, "images": changed} - finally: - connection.close() - - -def _cleanup(args: argparse.Namespace) -> None: - if args.apply and not args.offline: - raise SystemExit("cleanup --apply requires --offline") - connection = _connect(Path(args.database).resolve()) - media = Path(args.media).resolve() - try: - connection.execute("BEGIN IMMEDIATE") - referenced: set[str] = set() - for row in connection.execute("SELECT content FROM conversations"): - try: - content = json.loads(row[0]) if isinstance(row[0], str) else row[0] - if not isinstance(content, list): - raise ValueError("history content is not a list") - for item in content: - if not isinstance(item, dict): - raise ValueError("history message is not an object") - pending = list(content) - while pending: - part = pending.pop() - if isinstance(part, list): - pending.extend(part) - elif isinstance(part, dict): - if part.get("type") == "image_media_ref": - parsed = ImageMediaRef( - media_id=part["media_id"], - mime_type=part["mime_type"], - width=part.get("width"), - height=part.get("height"), - byte_size=part["byte_size"], - detail=part.get("detail"), - version=part.get("version", 1), - image_id=part.get("image_id"), - ) - ImageMediaStore(media).read(parsed, {parsed.media_id}) - referenced.add(parsed.media_id) - else: - pending.extend(part.values()) - except Exception as exc: # noqa: BLE001 - raise RuntimeError("history parse failed; cleanup aborted") from exc - quarantine = Path( - args.quarantine or (media.parent / (media.name + ".quarantine")) - ) - if args.apply: - quarantine.mkdir(parents=True, exist_ok=True) - moved = 0 - for media_id in sorted( - {path.stem for path in media.iterdir() if path.suffix in {".bin", ".json"}} - ): - paths = [media / f"{media_id}.bin", media / f"{media_id}.json"] - if any(path.is_symlink() for path in paths if path.exists()): - continue - if media_id in referenced: - continue - existing_paths = [path for path in paths if path.exists()] - if len(existing_paths) != 2: - continue - try: - metadata = json.loads(paths[1].read_text()) - if ( - metadata.get("media_id") != media_id - or _digest(paths[0]) != media_id - ): - continue - except (OSError, json.JSONDecodeError): - continue - if args.apply: - targets = [quarantine / path.name for path in existing_paths] - if any(target.exists() for target in targets): - raise RuntimeError( - "quarantine target already exists; refusing overwrite" - ) - for path, target in zip(existing_paths, targets): - shutil.move(str(path), target) - moved += len(existing_paths) - if args.apply: - connection.commit() - else: - connection.rollback() - print( - json.dumps( - { - "dry_run": not args.apply, - "quarantined": moved, - "referenced": len(referenced), - } - ) - ) - finally: - connection.close() - - -def _restore(args: argparse.Namespace) -> None: - """Restore only when the database and media are unchanged since migration.""" - if not args.offline: - raise SystemExit("restore requires --offline") - backup = Path(args.restore).resolve() - db = Path(args.database).resolve() - backup_db = backup / db.name - manifest_path = backup / "manifest.json" - if not backup_db.is_file() or not manifest_path.is_file(): - raise SystemExit("backup is incomplete") - manifest = json.loads(manifest_path.read_text()) - if _digest(backup_db) != manifest.get("database"): - raise SystemExit("restore refused: archived database hash mismatch") - archived_media = backup / "media" - manifest_names: set[str] = set() - for item in manifest.get("media", []): - _manifest_media_name(item["name"]) - checksum = item.get("sha256", "") - if ( - item["name"] in manifest_names - or len(checksum) != 64 - or any(char not in "0123456789abcdef" for char in checksum) - ): - raise SystemExit("restore refused: invalid media manifest") - manifest_names.add(item["name"]) - source = archived_media / item["name"] - if ( - source.is_symlink() - or not source.is_file() - or _digest(source) != item["sha256"] - ): - raise SystemExit("restore refused: archived media hash mismatch") - current = _connect(db) - try: - if _logical_digest(current) != manifest.get("post_migration_logical_digest"): - raise SystemExit("restore refused: database changed since migration") - finally: - current.close() - media = Path(args.media).resolve() - for item in manifest.get("media", []): - target = media / item["name"] - if target.exists() and ( - target.is_symlink() or _digest(target) != item["sha256"] - ): - raise SystemExit(f"restore refused: conflicting media: {target.name}") - rollback = db.with_name(db.name + ".before-restore") - if rollback.exists(): - raise SystemExit("refusing to overwrite an existing pre-restore copy") - try: - staged_media: list[tuple[Path, Path]] = [] - staging = Path( - tempfile.mkdtemp(prefix="astrbot-restore-media-", dir=media.parent) - ) - for item in manifest.get("media", []): - source = archived_media / item["name"] - staged = staging / item["name"] - shutil.copy2(source, staged) - if _digest(staged) != item["sha256"]: - raise SystemExit("restore refused: staged media hash mismatch") - staged_media.append((staged, media / item["name"])) - media.mkdir(parents=True, exist_ok=True) - for staged, target in staged_media: - if not target.exists(): - staged.replace(target) - - rollback_connection = sqlite3.connect(rollback) - current = _connect(db) - try: - current.backup(rollback_connection) - rollback_connection.commit() - finally: - current.close() - rollback_connection.close() - destination = _connect(db) - archived = _connect(backup_db) - try: - archived.backup(destination) - destination.commit() - finally: - archived.close() - destination.close() - finally: - if "staging" in locals(): - shutil.rmtree(staging, ignore_errors=True) - print(json.dumps({"restored": str(db), "pre_restore_copy": str(rollback)})) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("database") - parser.add_argument("--media", required=True) - parser.add_argument("--conversation-id", action="append") - parser.add_argument("--apply", action="store_true") - parser.add_argument("--offline", action="store_true") - parser.add_argument("--backup") - parser.add_argument("--cleanup", action="store_true") - parser.add_argument("--quarantine") - parser.add_argument("--restore") - args = parser.parse_args() - if args.restore: - _restore(args) - elif args.cleanup: - _cleanup(args) - else: - _migrate(args) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/unit/test_image_lifecycle_workloads.py b/tests/unit/test_image_lifecycle_workloads.py deleted file mode 100644 index 42427d35e..000000000 --- a/tests/unit/test_image_lifecycle_workloads.py +++ /dev/null @@ -1,36 +0,0 @@ -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path - - -def test_lifecycle_workload_defines_bounded_matrix_and_no_payload_dump(): - source = Path("scripts/image_memory_bench/lifecycle_workloads.py").read_text( - encoding="utf-8" - ) - assert "REQUESTS = 200" in source - assert "SESSIONS = 4" in source - assert "--request-count" in source - assert "--window-turns" in source - assert "time.sleep(0.01)" in source - assert "120" in source - assert "b64encode" not in source - assert "print(" not in source - - -def test_lifecycle_child_requires_real_database_and_reports_wire_metadata(tmp_path): - output = tmp_path / "result.jsonl" - result = subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/lifecycle_workloads.py", - str(output), - "--child", - ], - check=False, - capture_output=True, - text=True, - ) - assert result.returncode != 0 - assert not output.exists() diff --git a/tests/unit/test_image_memory_benchmark.py b/tests/unit/test_image_memory_benchmark.py deleted file mode 100644 index 65ec479c9..000000000 --- a/tests/unit/test_image_memory_benchmark.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Safety checks for the standalone image benchmark utilities.""" - -import json -import subprocess -import sys - - -def test_baseline_manifest_records_fixture_fingerprints(tmp_path): - fixture = tmp_path / "fixtures" / "7" - fixture.mkdir(parents=True) - (fixture / "sample.png").write_bytes(b"fixture") - output = tmp_path / "manifest.json" - subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/baseline_manifest.py", - str(output), - "--fixtures", - str(tmp_path / "fixtures"), - ], - check=True, - ) - manifest = json.loads(output.read_text()) - assert manifest["fixtures"][0]["sha256"] - - -def test_migration_dry_run_and_rollback_keep_input(tmp_path): - source = tmp_path / "history.jsonl" - source.write_text( - json.dumps({"history": [{"role": "user", "content": "hello"}]}) + "\n" - ) - output = tmp_path / "migrated.jsonl" - media = tmp_path / "media" - command = [ - sys.executable, - "scripts/image_memory_bench/migrate_history.py", - str(source), - str(output), - "--media-dir", - str(media), - ] - subprocess.run(command, check=True) - assert not output.exists() - subprocess.run(command + ["--apply"], check=True) - rollback = tmp_path / "rollback.jsonl" - subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/migrate_history.py", - str(output), - str(rollback), - "--media-dir", - str(media), - "--rollback", - "--apply", - ], - check=True, - ) - assert output.read_bytes() == rollback.read_bytes() diff --git a/tests/unit/test_manage_image_history.py b/tests/unit/test_manage_image_history.py deleted file mode 100644 index 959db2f52..000000000 --- a/tests/unit/test_manage_image_history.py +++ /dev/null @@ -1,354 +0,0 @@ -from __future__ import annotations - -import base64 -import importlib.util -import json -import sqlite3 -import subprocess -import sys - -import pytest -from PIL import Image - - -def _image() -> bytes: - import io - - output = io.BytesIO() - Image.new("RGB", (3, 2), "red").save(output, "PNG") - return output.getvalue() - - -def _db(path, content): - connection = sqlite3.connect(path) - connection.execute( - "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" - ) - connection.execute( - "INSERT INTO conversations VALUES (1, 'one', ?)", (json.dumps(content),) - ) - connection.commit() - connection.close() - - -def _run(db, media, *args): - return subprocess.run( - [ - sys.executable, - "scripts/manage_image_history.py", - str(db), - "--media", - str(media), - *args, - ], - check=False, - capture_output=True, - text=True, - ) - - -def test_migration_defaults_to_dry_run(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode() - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, history) - result = _run(db, media) - assert result.returncode == 0 - assert not media.exists() - assert ( - json.loads( - sqlite3.connect(db) - .execute("SELECT content FROM conversations") - .fetchone()[0] - ) - == history - ) - - -def test_apply_creates_backup_and_references(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode(), - "id": "keep", - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) - assert result.returncode == 0, result.stderr - content = json.loads( - sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] - ) - ref = content[0]["content"][0] - assert ref["type"] == "image_media_ref" and ref["image_id"] == "keep" - assert (backup / "db.sqlite").exists() and (backup / "manifest.json").exists() - - -def test_apply_requires_offline_and_cleanup_dry_run_is_read_only(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, []) - media.mkdir() - (media / "dead.bin").write_bytes(b"x") - result = _run(db, media, "--apply") - assert result.returncode != 0 - result = _run(db, media, "--cleanup") - assert result.returncode == 0 - assert (media / "dead.bin").exists() - - -def test_cleanup_aborts_on_malformed_history(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, {"not": "a history"}) - media.mkdir() - (media / ("a" * 64 + ".bin")).write_bytes(b"orphan") - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode != 0 - assert (media / ("a" * 64 + ".bin")).exists() - - -def test_cleanup_preserves_references_nested_in_retained_messages(tmp_path): - from astrbot.core.utils.image_media_store import ImageMediaStore - - media = tmp_path / "media" - store = ImageMediaStore(media) - ref = store.put(_image()) - db = tmp_path / "db.sqlite" - _db(db, [{"role": "tool", "content": [{"resource": ref.model_dump()}]}]) - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode == 0, result.stderr - assert store.read(ref, {ref.media_id}) == _image() - - -def test_cleanup_keeps_shared_media_until_last_history_reference_is_removed(tmp_path): - from astrbot.core.utils.image_media_store import ImageMediaStore - - media = tmp_path / "media" - store = ImageMediaStore(media) - ref = store.put(_image()) - first_history = [{"role": "user", "content": [ref.model_dump()]}] - second_history = [{"role": "user", "content": [ref.model_dump()]}] - db = tmp_path / "db.sqlite" - _db(db, first_history) - connection = sqlite3.connect(db) - connection.execute( - "INSERT INTO conversations VALUES (2, 'two', ?)", - (json.dumps(second_history),), - ) - connection.commit() - connection.close() - - first_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert first_delete.returncode == 0, first_delete.stderr - assert (media / f"{ref.media_id}.bin").exists() - - connection = sqlite3.connect(db) - connection.execute("DELETE FROM conversations WHERE inner_conversation_id=1") - connection.commit() - connection.close() - second_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert second_delete.returncode == 0, second_delete.stderr - assert (media / f"{ref.media_id}.bin").exists() - - connection = sqlite3.connect(db) - connection.execute("DELETE FROM conversations WHERE inner_conversation_id=2") - connection.commit() - connection.close() - last_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert last_delete.returncode == 0, last_delete.stderr - assert not (media / f"{ref.media_id}.bin").exists() - quarantine = media.parent / "media.quarantine" - assert (quarantine / f"{ref.media_id}.bin").exists() - - -def test_cleanup_leaves_unpaired_and_symlink_objects(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, []) - media.mkdir() - media_id = "b" * 64 - (media / f"{media_id}.bin").write_bytes(b"bad") - (media / f"{media_id}.json").write_text("{}") - (media / ("c" * 64 + ".bin")).write_bytes(b"bad") - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode == 0 - assert (media / f"{media_id}.bin").exists() - - -def test_restore_refuses_newer_history_and_wal_is_not_left_stale(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode() - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - connection = sqlite3.connect(db) - connection.execute("UPDATE conversations SET content='[]'") - connection.execute("CREATE TABLE unrelated (value TEXT)") - connection.execute("INSERT INTO unrelated VALUES ('changed')") - connection.commit() - connection.close() - (db.with_name(db.name + "-wal")).write_bytes(b"stale") - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode != 0 - assert not (db.with_name(db.name + ".before-restore")).exists() - - -def test_restore_validates_backup_and_conflicting_media_before_database_write(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, []) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - manifest = backup / "manifest.json" - manifest.write_text( - manifest.read_text().replace( - '"database":', '"database": "tampered", "ignored":' - ) - ) - before = db.read_bytes() - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode != 0 - assert db.read_bytes() == before - - -def test_successful_restore_preserves_inline_history_and_detail(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode(), - "detail": "high", - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode == 0, result.stderr - restored = json.loads( - sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] - ) - assert restored == history - - -def test_backup_reads_consistent_open_wal_database(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - connection = sqlite3.connect(db) - connection.execute("PRAGMA journal_mode=WAL") - connection.execute( - "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" - ) - connection.execute("INSERT INTO conversations VALUES (1, 'wal', '[]')") - connection.commit() - assert (db.with_name(db.name + "-wal")).exists() - connection.close() - result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) - assert result.returncode == 0, result.stderr - assert ( - sqlite3.connect(backup / "db.sqlite") - .execute("SELECT conversation_id FROM conversations") - .fetchone()[0] - == "wal" - ) - - -def test_restore_replace_failure_leaves_database_and_cleans_stage( - tmp_path, monkeypatch -): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, []) - from astrbot.core.utils.image_media_store import ImageMediaStore - - ImageMediaStore(media).put(_image(), detail="high") - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - for path in media.iterdir(): - path.unlink() - module_spec = importlib.util.spec_from_file_location( - "manage_image_history", "scripts/manage_image_history.py" - ) - module = importlib.util.module_from_spec(module_spec) - module_spec.loader.exec_module(module) - - def fail_replace(self, target): - raise OSError("injected replace failure") - - monkeypatch.setattr(module.Path, "replace", fail_replace) - args = type( - "Args", - (), - { - "offline": True, - "restore": str(backup), - "database": str(db), - "media": str(media), - }, - ) - before = db.read_bytes() - with pytest.raises(OSError, match="injected replace failure"): - module._restore(args) - assert db.read_bytes() == before - assert not list(tmp_path.glob("astrbot-restore-media-*"))