mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-10-16 15:06:04 +08:00
test: add image memory benchmarks and cross-platform validation
This commit is contained in:
zenfun
@@ -0,0 +1,91 @@
|
||||
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: 45
|
||||
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: |
|
||||
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_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_provider_request_retry.py \
|
||||
tests/unit/test_manage_image_history.py \
|
||||
tests/unit/test_conversation_media_api.py
|
||||
|
||||
- name: Generate reproducible fixtures outside the measured process
|
||||
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
|
||||
|
||||
- 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 ordinary and stress memory harness workloads
|
||||
run: |
|
||||
result="$RUNNER_TEMP/image-memory-validation/results/measurements.jsonl"
|
||||
while IFS= read -r image; do
|
||||
uv run python scripts/image_memory_bench/measure.py \
|
||||
"$image" "$result" \
|
||||
--repo "$GITHUB_WORKSPACE" \
|
||||
--workload "$(basename "$image")"
|
||||
done < <(find "$RUNNER_TEMP/image-memory-validation/fixtures/7" -type f \
|
||||
\( -name '*ordinary*' -o -name '*stress*' \) | sort)
|
||||
|
||||
- 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
|
||||
@@ -0,0 +1,2 @@
|
||||
schema: spec-driven
|
||||
created: 2026-09-15
|
||||
@@ -0,0 +1,152 @@
|
||||
## Context
|
||||
|
||||
See `proposal.md` and the three delta specs. The current provider request path turns image references into complete data URIs, while conversation saving serializes the model-visible message list. Image preparation currently focuses on longest-edge pixels and can pass through large pixel-compliant files. Existing histories already contain inline data URIs and must remain readable.
|
||||
|
||||
## Goals / Non-Goals
|
||||
|
||||
**Goals:**
|
||||
|
||||
- Establish one preparation boundary shared by all image-producing inputs.
|
||||
- Bound the encoded payload that can enter a provider request.
|
||||
- Separate durable conversation records from media bytes and resolve media only for selected active messages.
|
||||
- Keep old inline histories readable and make migration explicit.
|
||||
- Measure file size, encoded payload size, request size, RSS high-water mark, Python allocations, temporary files, and post-request retention.
|
||||
|
||||
**Non-Goals:**
|
||||
|
||||
- Do not silently remove images from the active context merely to improve memory metrics.
|
||||
- Do not rewrite existing conversations during ordinary reads.
|
||||
- Do not assume the 5 MiB limit from #10089 applies to every provider.
|
||||
- Do not make the provider-specific wire format depend on the persistence representation.
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. Use a media reference as the persistence boundary
|
||||
|
||||
New history image parts use a versioned internal media reference containing a content hash, MIME type, dimensions, encoded byte size, and original image-detail metadata. Durable image objects live in a dedicated media directory under the configured AstrBot data root, NEVER the temporary directory or its age-based cleanup policy. Provider adapters continue to receive their existing image content shape after request-time materialization. No extra image-caption model is introduced.
|
||||
|
||||
Alternative rejected: keep base64 in history and only trim it during compaction. This reduces some requests but leaves database growth, repeated JSON parsing, and compaction-record duplication.
|
||||
|
||||
### 2. Keep backward-compatible dual readers
|
||||
|
||||
The history reader accepts both the new reference form and existing inline `data:`/base64 forms. The writer emits references only for newly prepared images. Conversion of existing records is an explicit repair/migration operation with backup, dry-run statistics, hash verification, and rollback by restoring the original history.
|
||||
|
||||
Alternative rejected: automatic in-place migration on first read. It makes a read destructive and risks losing recoverability when the media directory is unavailable.
|
||||
|
||||
### 3. Make encoded bytes a first-class limit
|
||||
|
||||
Image preparation uses independent limits for maximum dimensions and maximum encoded payload bytes. Compute base64 length as `4 * ceil(encoded_file_bytes / 3)` without allocating a base64 copy for every candidate. Retain only the current best candidate. Preserve already-compliant original bytes; otherwise try the original supported format and JPEG quality steps, select a valid candidate without enlarging compliant source content, and encode base64 only at the provider boundary. Preserve transparency when required. If no candidate fits, return a readable size error rather than silently returning the original oversized image.
|
||||
|
||||
Proposed product defaults: preserve existing dimension and quality settings, add `image_compress_options.max_encoded_bytes = 4194304` (4 MiB of base64 per image). This is an AstrBot preparation budget, not a claim about vendor limits. A documented or configured smaller provider limit takes precedence. A known aggregate request limit is checked against the complete serialized request, independently of this per-image budget. Never reinterpret base64 bytes as text tokens. No universal aggregate vendor limit is invented.
|
||||
|
||||
For CUA, preserve the oriented original pixel dimensions and coordinates; try encoding changes without resizing and return a readable error if the budget cannot be met. Do not use an enormous pixel limit as a substitute for an explicit preserve-dimensions flag. Animation follows the selected baseline's frame-selection behavior; an older branch that skips animations is not silently treated as the current montage implementation. Do not keep all decoded animation frames resident simultaneously.
|
||||
|
||||
Alternative rejected: use only pixel dimensions or only source-file bytes. Neither predicts the final JSON payload reliably.
|
||||
|
||||
### 4. Materialize only selected references
|
||||
|
||||
Context selection happens before resolving media references. Images not selected for either the main request or a separate image-capable summary request are not opened, decoded, or encoded. The summary request and the subsequent main request have separate measured lifetimes. Selected references are materialized into a request-local structure and released after the request; the history model never receives the resulting data URI. In-flight cancellation must not delete a file still used by a non-cancelled image worker; cleanup runs when that worker actually exits.
|
||||
|
||||
For a fixed provider configuration, prepared historical bytes are immutable. Do not re-encode the history when a new image arrives, tune old JPEG quality to the remaining request space, or substitute captions to achieve memory targets. Ordinary summary/truncation semantics are unchanged. A known oversized aggregate request fails clearly; this project does not add image eviction or a new compression trigger. B-only lazy loading MUST send byte-identical image content to the baseline.
|
||||
|
||||
Alternative rejected: materialize all references first and let the context manager remove messages afterward. That preserves the current memory spike and defeats lazy loading.
|
||||
|
||||
### 5. Treat all sources as adapters to one preparation function
|
||||
|
||||
Platform attachments, quoted content, plugin/MCP results, file tools, and CUA screenshots normalize to a common input descriptor. The descriptor carries source ownership and cleanup responsibility. The preparation function owns format detection, dimension checks, byte checks, temporary-file cleanup, and diagnostic metadata.
|
||||
|
||||
### 6. Use an external-process ablation harness
|
||||
|
||||
Each measurement case runs in a fresh child process. The harness records RSS at high frequency, Python allocation snapshots separately, request JSON size, image byte sizes, media object counts, and temporary files. Cold-start, warm-start, single-image, multi-image, long-history, concurrent, restart, and missing-media cases are separate workloads.
|
||||
|
||||
The implementation is accepted only when the optimized path preserves image count/order and provider-visible content for the active window. Memory reduction caused by silently dropping images is a failed experiment, not a success.
|
||||
|
||||
## Risks / Trade-offs
|
||||
|
||||
- [Risk] Media references can outlive their files. → Keep content hashes and metadata, report a bounded placeholder on misses, and add cleanup/reconciliation diagnostics.
|
||||
- [Risk] Existing plugins may assume persisted `image_url` always contains a data URI. → Keep the dual reader and resolve references at the provider boundary; add compatibility tests for plugin/tool history.
|
||||
- [Risk] Extra disk I/O can increase latency. → Resolve only selected images, cache prepared media by content hash within one request, and measure cold/warm latency separately.
|
||||
- [Risk] Image conversion can use native memory invisible to `tracemalloc`. → Use an external RSS monitor and report both RSS and Python allocations.
|
||||
- [Risk] A provider may have a smaller limit than the internal default. → Apply the effective provider limit during request preparation and preserve readable provider-specific errors.
|
||||
|
||||
## Migration Plan
|
||||
|
||||
1. Ship dual reading before new reference writing; factor switches belong in the experiment harness, not a new public rollout switch.
|
||||
2. Run the ablation suite and compatibility tests before enabling reference writing by default.
|
||||
3. Provide a dry-run migration report for old inline histories.
|
||||
4. Migrate selected conversations only after media files and hashes are verified.
|
||||
5. Before downgrading to a reader without reference support, stop writes and export/re-inline all reference-bearing histories into a verified backup, including conversations created after deployment. Merely disabling new writes does not make those records readable by old versions. Retain the media store until rollback verification completes.
|
||||
|
||||
## 兼容与资源生命周期约束
|
||||
|
||||
- 当前基线工作区提交为 `1a2492a09f8722a9abd813f6a65fb687d53c88c9`，位于此前功能分支，不等同于 issue 的 v4.28.1。实现前固定目标基线及依赖，另建 v4.28.1 只读复现环境；不能拿两个版本的路径混合计算收益。
|
||||
- 新引用采用内部版本化图片部件，保留 image detail；对外 ProviderRequest 输入继续接受原来的路径、URL、data URI。新类型不得直接发送给远端 Provider。
|
||||
- Plugin/第三方 Provider 的兼容出口必须解析引用后再调用现有接口；内置 Provider 和摘要 Provider 都纳入测试。给插件的兼容视图可能仍需分配 base64，这部分成本如实报告，不能声称所有消费者都已懒加载。
|
||||
- 媒体文件采用内容哈希去重，先在同文件系统原子写入并核验，再提交消息引用。写入失败保持原会话不变；崩溃留下的未引用文件由维护命令处理。
|
||||
- 媒体读取通过内部 ID 查询并限制路径；WebUI 使用带会话权限校验的图片端点，不公开数据目录或仅凭哈希授权。会话列表/详情只返回引用和预览地址，不批量还原 base64。
|
||||
- 复用现有附件存储能力前核对其 WebChat 删除逻辑，禁止把 Agent 媒体直接挂入会被其他界面独立删除的生命周期。引用对象与消息内容绑定，跨会话去重不得带来跨会话读权限。
|
||||
- 首版不做运行中自动垃圾回收。显式维护命令在停止写入后扫描全部保留会话，输出 dry-run，隔离未引用对象；读写中断、解析失败时停止删除。移除或压缩一个会话不能删除仍被其他会话引用的文件。
|
||||
- 旧 inline 图片在普通读取时不转码、不改写；显式迁移只搬运原字节并核验哈希，不在迁移中夹带画质调整。历史包与媒体目录共同备份；导出保持自包含或附带媒体清单，导入必须校验引用。
|
||||
- 本地资源异常不应被笼统吞掉再走同一大请求；测试检查 MemoryError 保留异常类型并终止本次运行。413 与资源耗尽分开记录，不假设每个回退 Provider 的上限都相同，不引入自动多轮重压缩循环。
|
||||
- 这不是对话事件日志重构；当前正常上下文压缩及保存行为保持不变。媒体表示迁移与“保留所有原始对话”的产品需求不得混为一谈。
|
||||
|
||||
## 详细内存消融实验
|
||||
|
||||
### 因素与对照
|
||||
|
||||
三个开关仅供实验调用真实生产函数，不提供给用户。完整八组为 `000 / A00 / 0B0 / 00C / AB0 / A0C / 0BC / ABC`。
|
||||
|
||||
| 因素 | 唯一变化 | 必须保持不变 |
|
||||
| --- | --- | --- |
|
||||
| A 统一入口 | 所有来源进入同一准备链 | 使用原压缩算法和参数 |
|
||||
| B 懒加载 | 同一准备结果外置存储，选窗后加载 | 图片字节、MIME、顺序、detail、消息位置 |
|
||||
| C 压缩优化 | 仅在基线已有入口改变压缩算法 | 不顺便补齐工具/插件入口 |
|
||||
|
||||
八组的零开关路径必须与固定基线的输入输出一致；用独立基线 checkout 的相同函数校验，不能把一个近似模拟器叫基线。另加纯文本对照，但不把它算作三因素收益。
|
||||
|
||||
### 固定样本
|
||||
|
||||
五类图片各包含普通与压力样本，每类三个固定种子，至少 30 个实例。文件在测量进程外生成，记录 SHA256、编码格式、尺寸、方向、透明度、帧数和字节数。
|
||||
|
||||
| 类别 | 普通 | 压力 |
|
||||
| --- | --- | --- |
|
||||
| PNG | 透明图/色块 | 1280×1280 高熵 RGBA，像素合规但字节超限 |
|
||||
| JPEG | 常见照片、已压缩小文件 | 4000×3000 高细节、EXIF 旋转 |
|
||||
| WebP | 静态有损 | 无损透明，动画分支另做覆盖 |
|
||||
| GIF | 少帧动画 | 多帧动画，检查帧抽取/拼图峰值 |
|
||||
| 工具截图 | 1920×1080 终端 | 4K 小字 UI 与坐标标记 |
|
||||
|
||||
来源测试覆盖本地文件、HTTP、data URI、base64 URI，以及用户附件、引用、插件/MCP、FileRead、CUA。HTTP 使用本地服务器；这些入口调用现有适配代码，不能仅给同一函数贴不同来源名称。完整交叉用于单图工作负载，其余长会话用五种格式均衡混合。
|
||||
|
||||
### 工作负载
|
||||
|
||||
1. 单图单轮、同轮八张不同图片；分别统计每个阶段。
|
||||
2. 50 轮每轮一张不同图片，然后 20 轮纯文本；与重复同一张图片的独立测试区分去重收益。
|
||||
3. 10/50/100 张历史图；分别测全在窗口、按轮次裁剪、摘要成功和摘要失败四种状态。记录实际触发次数，不假设默认配置一定触发。
|
||||
4. 摘要模型分别支持和不支持图片。固定摘要返回值及用量，避免随机文本干扰；摘要图像输入遵循实际模态规则。
|
||||
5. 四会话并发、保存后重启、WebUI 只看详情/展开单张预览、显式导出；WebUI 与模型请求分别统计。
|
||||
6. 错误实验：缺失媒体、坏图片、取消、超时、写盘失败、5 MiB 聚合请求拒绝、独立单图限制、不同限制的回退 Provider。
|
||||
7. 静态图片循环固定窗口请求 200 次，记录 1/10/25/50/100/200 次，检查存活对象、临时文件、打开句柄及运行后残留；不强制 GC 的正常数据与诊断 GC 分开。
|
||||
|
||||
### 测量与防止实验污染
|
||||
|
||||
- 分阶段打点：读取数据库 → JSON 解析/消息构造 → 选窗 → 读取媒体 → 解码/缩放/编码 → base64 → Provider 组装 → HTTP 序列化/发送 → 保存 → 清理。另标记摘要请求阶段。
|
||||
- 记录原文件与输出文件字节、base64 长度、完整 HTTP JSON 字节、读取/解码/转码/base64 次数及字节量、阶段耗时、临时文件和持久化字节。磁盘数据不等价于内存收益。
|
||||
- 外部监测进程每 10 ms 采集被测进程及子进程 RSS；记录 OS 高水位。Windows 另采私有内存与工作集。两个平台单独比较，不直接比较绝对值。
|
||||
- 标准轮关闭 tracemalloc；诊断轮单独启用 tracemalloc，Linux 关键图另外追踪原生分配。先前约 24 MB WebP 结果缺乏初始化隔离，不能用作验收基线。
|
||||
- 每个实验单元十次独立进程运行，冷启动/预热三次后的热运行分别记录。按种子配对，随机交错八组，保存完整运行顺序。固定机器、解释器、依赖和处理并发数。
|
||||
- 本地模型服务返回固定文本、固定用量和固定摘要；真实客户端/SDK 必须参与最终请求组装。服务端流式计数，不保留所有历史请求体；结构校验另跑，不污染测量峰值。
|
||||
- 监测器不拷贝 base64、不输出原图、会话正文或密钥。样本生成、结果分析和监测服务不得计入被测进程。
|
||||
- 单次 120 秒或被测进程总内存达到 `min(2 GiB, 启动时可用内存的 25%)` 时停止并记录超限。崩溃、超时、OOM 数据不得丢弃或用零内存填充。
|
||||
|
||||
### 分析和验收
|
||||
|
||||
- 单因素效果比较 A/B/C 与基线，边际效果比较 ABC 与 BC/AC/AB；报告配对中位差、范围与原始样本，不把三个百分比相加。十次重复不宣称精确的 P99。
|
||||
- B-only 校验图片哈希、顺序、detail 和固定 Provider 的历史前缀；预期最终请求大小不变。请求时仍需装载整个有效图片窗口，必须披露这个下限。
|
||||
- C 校验方向、透明、截图文字和坐标；已合规原字节应原样保留，不能选用无必要膨胀的输出。超限失败不算“成功压缩”或内存收益。
|
||||
- CUA 的尺寸和坐标必须一致；失败不能伪装成空图成功。GIF 的既有抽帧语义必须一致。
|
||||
- B 在 50 张历史图的读取到选窗阶段，以相同消息窗口为前提，峰值增量降低至少 50%；ABC 在高熵长会话的成功可比场景中，整体峰值增量降低至少 30%。这是验收目标，不是已证明结果。
|
||||
- 小图正常场景峰值增量及中位耗时回退不超过 10%；图片质量或稳定历史前缀失败时，即使内存达标也不通过。
|
||||
- 持续请求后存活图片对象、文件句柄、临时文件不能按轮数线性增长。RSS 分配器保留不直接判定泄漏；结合对象与原生分配数据解释。
|
||||
- 实测平台最低覆盖 Linux/WSL 与原生 Windows，macOS 做功能回归。所有指标原始 JSONL/CSV、样本清单和曲线放在测试产物目录，不创建仓库 SUMMARY 报告。
|
||||
- 真实智谱只用于可选小样本画质/协议验证，与消融数据隔离。默认不发公网请求；密钥通过安全环境注入，绝不写入本规划或测试产物。
|
||||
@@ -0,0 +1,28 @@
|
||||
## Why
|
||||
|
||||
AstrBot currently prepares images through several partially independent paths. Some inputs are resized only by pixel dimensions, while the final base64 payload can still exceed provider limits and consume several copies of the image in memory. Conversation history can also retain complete inline image data, so every later request pays the deserialization and request-construction cost again. Issue #10089 demonstrates request-size failures, while #10092 demonstrates process-level memory exhaustion.
|
||||
|
||||
## What Changes
|
||||
|
||||
- Add one image preparation contract for user attachments, quoted images, plugins, MCP/tool results, file reading, and screenshots.
|
||||
- Enforce both dimension limits and a configurable final encoded-payload byte limit; reject or degrade images that cannot satisfy the limit instead of silently sending oversized data.
|
||||
- Store historical images as durable media references and metadata rather than embedding complete base64 data in every conversation message.
|
||||
- Resolve historical image references only when the active provider request needs them, with bounded loading and cleanup.
|
||||
- Keep the current conversation semantics: images in the active context remain available to the model; this change does not silently remove images merely to improve memory numbers.
|
||||
- Add memory-ablation tests covering preparation, history loading, request assembly, persistence, and repeated conversations.
|
||||
- Preserve backward compatibility by reading existing inline data URIs and provide an explicit migration/repair path instead of rewriting them implicitly.
|
||||
|
||||
## Capabilities
|
||||
|
||||
### New Capabilities
|
||||
|
||||
- `image-memory-lifecycle`: bounded image preparation, durable media references, lazy resolution, and memory-safe request assembly.
|
||||
- `agent-context-image-budget`: active requests validate image byte size independently from token accounting and resolve references without changing image visibility.
|
||||
- `conversation-history-media`: persisted messages may use media references and must remain readable across restart while preserving existing inline-image histories.
|
||||
|
||||
## Impact
|
||||
|
||||
- Affects media resolution and compression, provider request assembly, agent context processing, conversation persistence, tool/plugin image paths, and related tests.
|
||||
- Adds a local media storage/index lifecycle with cleanup and missing-media handling.
|
||||
- Adds configuration for final encoded image size and bounded media storage behavior; existing pixel/quality settings remain compatible.
|
||||
- Provider payloads remain provider-specific; only the internal history representation and preparation boundary change.
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
## Purpose
|
||||
|
||||
Keeps the active model context bounded by treating image payload bytes as a real request cost and by resolving only image data required by the active provider request.
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: Active context SHALL use a bounded image budget
|
||||
The system SHALL validate encoded image bytes and known provider request-size limits independently from text token estimates. It MUST NOT introduce additional image eviction, replace available images with summaries, or re-encode historical images on every turn to fit a changing budget.
|
||||
|
||||
#### Scenario: Context contains many images
|
||||
- **WHEN** the active context contains images whose encoded payloads exceed the image budget
|
||||
- **THEN** the system reports a readable size-limit error before sending a known-oversized request, without silently dropping images or changing historical image bytes
|
||||
|
||||
### Requirement: Historical media SHALL be resolved on demand
|
||||
The system SHALL keep historical image content as a resolvable media reference and SHALL materialize its bytes only when the image is selected for the active provider request.
|
||||
|
||||
#### Scenario: Old image is outside the active context
|
||||
- **WHEN** a historical image is outside the selected context window and is not input to the summarization request
|
||||
- **THEN** the system does not read, decode, or base64-encode that image for the request
|
||||
|
||||
### Requirement: Summarization SHALL retain its existing image input semantics
|
||||
The system SHALL resolve selected image references for an image-capable summarization provider just as it does for the main provider. Images summarized out of the main window may still need loading for that separate summary request; this cost SHALL be measured separately.
|
||||
|
||||
#### Scenario: Old image participates in a summary
|
||||
- **WHEN** the existing context compression policy passes an old image to an image-capable summary provider
|
||||
- **THEN** that provider receives the image rather than an unresolved reference, while the subsequent main request contains only its selected context
|
||||
|
||||
### Requirement: Lazy loading SHALL preserve provider-visible image bytes
|
||||
The system SHALL preserve image bytes, MIME type, detail, order, and message placement for a fixed prepared image across persistence, restart, and subsequent requests to the same provider configuration.
|
||||
|
||||
#### Scenario: A later turn loads a persisted image
|
||||
- **WHEN** a later text message causes an existing image reference to be materialized
|
||||
- **THEN** the image content is byte-identical to its earlier prepared representation and the stable historical request prefix remains unchanged
|
||||
|
||||
### Requirement: Existing inline images SHALL remain readable
|
||||
The system SHALL continue to read existing persisted inline image data and SHALL apply the active image budget when such data is included in a request.
|
||||
|
||||
#### Scenario: Existing conversation contains a data URI
|
||||
- **WHEN** a conversation created before media references is loaded
|
||||
- **THEN** the conversation remains readable and its inline image is handled by the same bounded request policy
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
## Purpose
|
||||
|
||||
Separates durable conversation facts from large media payloads so that conversations remain restartable and image-capable without duplicating complete base64 data in every historical message.
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: New persisted image content SHALL use media references
|
||||
The system SHALL persist a stable media reference with MIME type and image metadata for newly saved model-visible images instead of embedding the complete base64 payload in the conversation record.
|
||||
|
||||
#### Scenario: Model-visible image is saved
|
||||
- **WHEN** a request containing a newly prepared image is saved to conversation history
|
||||
- **THEN** the history stores a media reference and sufficient metadata to resolve the image later
|
||||
|
||||
### Requirement: Media references SHALL survive restart
|
||||
The system SHALL resolve persisted media references after restart and SHALL report a readable missing-media result when the referenced media is unavailable.
|
||||
|
||||
#### Scenario: Referenced media file is missing
|
||||
- **WHEN** a conversation contains a reference whose media object no longer exists
|
||||
- **THEN** loading the conversation does not crash, the UI reports unavailable media, and the model receives an explicit bounded missing-image placeholder; this is reported as a recovery condition, not a successful memory optimization
|
||||
|
||||
### Requirement: Media lifetime SHALL be independent of temporary cleanup
|
||||
Referenced media SHALL be durable outside the temporary-file cleanup domain. The system SHALL authorize image reads through conversation access, preserve shared media until no retained history references it, and never interpret a client-controlled media reference as an arbitrary local path.
|
||||
|
||||
#### Scenario: Temporary files are cleaned after restart
|
||||
- **WHEN** temporary media cleanup runs while a saved conversation still references an image
|
||||
- **THEN** that saved image remains readable
|
||||
|
||||
#### Scenario: Another conversation requests an unauthorized image
|
||||
- **WHEN** a caller supplies an image identifier without access to its owning conversation
|
||||
- **THEN** the request is rejected without exposing the image or its filesystem path
|
||||
|
||||
### Requirement: History migration SHALL be explicit and reversible
|
||||
The system SHALL preserve existing inline-image histories and SHALL provide an explicit migration or repair operation before replacing inline payloads with media references.
|
||||
|
||||
#### Scenario: Existing history has inline base64
|
||||
- **WHEN** an old conversation is opened without migration
|
||||
- **THEN** it remains readable and no destructive rewrite occurs automatically
|
||||
@@ -0,0 +1,26 @@
|
||||
## Purpose
|
||||
|
||||
Provides one bounded lifecycle for images from users, tools, plugins, and external media so that model requests remain within provider limits without multiplying large image payloads in memory.
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: All model-bound images use one bounded preparation contract
|
||||
The system SHALL apply the same preparation contract to user attachments, quoted images, plugin and MCP results, file-reading images, and computer-use screenshots before they become model image content.
|
||||
|
||||
#### Scenario: Source-specific image enters a model request
|
||||
- **WHEN** an image is supplied by any supported source
|
||||
- **THEN** the system applies the same dimension, format, encoded-byte, cleanup, and failure rules before constructing provider content
|
||||
|
||||
### Requirement: Image preparation SHALL enforce final encoded size
|
||||
The system SHALL enforce a configurable upper bound on the final encoded image payload, in addition to pixel dimensions, and SHALL NOT silently send an image that exceeds the bound after preparation.
|
||||
|
||||
#### Scenario: Pixel-compliant image exceeds the byte limit
|
||||
- **WHEN** an image is within the configured dimensions but its encoded payload exceeds the configured byte limit
|
||||
- **THEN** the system re-encodes or resizes it until it fits, or returns a readable image-too-large failure without sending the oversized payload
|
||||
|
||||
### Requirement: Preparation SHALL release temporary resources
|
||||
The system SHALL clean resolver-owned temporary files and release image-processing resources after request construction succeeds, fails, is cancelled, or times out.
|
||||
|
||||
#### Scenario: Image preparation is cancelled
|
||||
- **WHEN** request processing is cancelled during download, decode, resize, or encoding
|
||||
- **THEN** temporary files and owned buffers are released without affecting unrelated media
|
||||
@@ -0,0 +1,36 @@
|
||||
## 1. Baseline and measurement harness
|
||||
|
||||
- [x] 1.1 Freeze the implementation and dependency baseline, create a separate read-only v4.28.1 reproduction environment for issues #10089 and #10092, and verify both environments report their commit and dependency versions.
|
||||
- [ ] 1.2 Add an external-process image benchmark harness that records stage timings, RSS high-water marks, Python allocation diagnostics, child-process memory, request JSON bytes, media bytes, temporary files, and failures without printing image data or secrets.
|
||||
- [x] 1.3 Generate deterministic PNG, JPEG, WebP, GIF, and screenshot fixtures with ordinary and high-entropy cases, record hashes and metadata, and verify the fixture generator is outside the measured child process.
|
||||
- [ ] 1.4 Implement the eight-factor ablation runner for unified-source preparation, lazy history loading, and compression optimization; verify paired seeds, cold/warm runs, ten repetitions, and raw JSONL/CSV output.
|
||||
|
||||
## 2. Unified preparation and encoded-byte limits
|
||||
|
||||
- [ ] 2.1 Define the internal image preparation descriptor and one preparation entry point for platform attachments, quoted images, plugins/MCP, file tools, and CUA screenshots; verify every listed source reaches it with source ownership and cleanup metadata.
|
||||
- [x] 2.2 Extend image preparation with an encoded payload byte limit independent from pixel dimensions; verify a pixel-compliant high-entropy PNG is re-encoded or rejected instead of passed through.
|
||||
- [ ] 2.3 Implement candidate-size measurement without retaining all candidate bytes and preserve transparency, EXIF orientation, CUA dimensions, animation behavior, and already-compliant bytes; verify format-specific fixtures and coordinate/visual checks.
|
||||
- [ ] 2.4 Add typed, readable image-size and resource errors and ensure oversize input is not silently returned to the Provider or retried through unrelated fallback providers; verify #10089's request-size scenario.
|
||||
- [x] 2.5 Add cancellation, timeout, decode failure, and write failure cleanup tests; verify temporary files, image-library resources, and worker tasks are released on every path.
|
||||
|
||||
## 3. Durable media references and lazy materialization
|
||||
|
||||
- [x] 3.1 Add versioned durable media objects under the configured data directory with content-hash deduplication, atomic write/verify, MIME/dimension/byte/detail metadata, and conversation-scoped access checks; verify partial writes never create usable references.
|
||||
- [ ] 3.2 Add a persisted image-reference representation and dual history reader for new references plus existing inline data URIs/base64; verify old histories remain readable and new histories do not embed complete base64.
|
||||
- [ ] 3.3 Resolve selected references only after context selection and materialize them into a request-local provider view; verify out-of-window images are not opened/decoded/encoded and B-only provider-visible bytes, order, detail, and placement match baseline.
|
||||
- [ ] 3.4 Keep durable media outside temporary cleanup and add missing-media, restart, unauthorized-reference, session-delete, and cross-session shared-media tests; verify missing media produces a bounded recovery result without exposing paths.
|
||||
- [ ] 3.5 Add explicit dry-run migration/export/repair behavior for old inline histories with backup and hash verification; verify ordinary reads never rewrite history and rollback remains possible.
|
||||
|
||||
## 4. Context and persistence integration
|
||||
|
||||
- [x] 4.1 Integrate image byte validation with the active request and summarization-provider paths without adding image eviction or changing existing truncation/summary semantics; verify image-capable summaries receive selected images and the main request receives its selected context.
|
||||
- [ ] 4.2 Ensure history saving persists references rather than request-time data URIs while provider adapters continue receiving their existing wire shapes; verify plugin, tool, built-in Provider, WebUI detail, and export compatibility.
|
||||
- [ ] 4.3 Separate 413 request-size, image-too-large, missing-media, and MemoryError handling; verify MemoryError preserves its type and does not trigger repeated oversized fallback requests.
|
||||
|
||||
## 5. Verification and rollout
|
||||
|
||||
- [ ] 5.1 Run the full matrix across single image, eight images, 50-round history, 200 fixed-window requests, four-session concurrency, restart, WebUI preview, and error workloads; verify all metrics and raw failures are retained.
|
||||
- [ ] 5.2 Compare A/B/C main and marginal effects with image-count/order/content invariants; verify memory gains are not caused by silently omitting active images.
|
||||
- [ ] 5.3 Validate Linux/WSL and native Windows memory behavior, run macOS functional regression, and verify normal small-image latency/peak-memory regression stays within the documented target.
|
||||
- [ ] 5.4 Run `ruff format --check .`, `ruff check .`, focused image/context/history tests, and the relevant full test suites; verify no source comments/logs violate repository language rules.
|
||||
- [ ] 5.5 Review migration backup/rollback and media cleanup behavior, document measured results in test artifacts outside repository SUMMARY files, and only then enable new reference writes by default.
|
||||
@@ -0,0 +1,54 @@
|
||||
"""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())
|
||||
@@ -0,0 +1,127 @@
|
||||
"""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())
|
||||
@@ -0,0 +1,361 @@
|
||||
"""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
|
||||
) -> 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)]
|
||||
# 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 media_root.name == "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("start")
|
||||
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("db_and_convert")
|
||||
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("json_load_bind")
|
||||
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")
|
||||
mark("selection")
|
||||
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, default=10)
|
||||
parser.add_argument("--warm-repeats", type=int, default=10)
|
||||
parser.add_argument("--warmup", type=int, default=3)
|
||||
parser.add_argument("--distinct", action="store_true")
|
||||
parser.add_argument("--csv", 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("--conversation-id", default="")
|
||||
parser.add_argument("--endpoint")
|
||||
parser.add_argument("--body-file", type=Path)
|
||||
parser.add_argument("--marker", type=Path)
|
||||
args = parser.parse_args()
|
||||
fixture_rows = json.loads(args.fixture.read_text())
|
||||
max_fixture_bytes = max((row.get("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:
|
||||
result = asyncio.run(run_child(args))
|
||||
metadata = json.loads(args.body_file.read_text())
|
||||
result.update(
|
||||
{
|
||||
"http_json_bytes": metadata["bytes"],
|
||||
"http_body_sha256": metadata["sha256"],
|
||||
"http_images": metadata["images"],
|
||||
}
|
||||
)
|
||||
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)
|
||||
csv_file = args.csv.open("w", newline="", encoding="utf-8") if args.csv else None
|
||||
writer = None
|
||||
received_body = {"path": None}
|
||||
|
||||
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(args.cold_repeats + args.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
|
||||
)
|
||||
)
|
||||
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 < args.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),
|
||||
"--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
|
||||
samples.append({"seconds": now, "rss": rss})
|
||||
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",
|
||||
}
|
||||
markers = [
|
||||
json.loads(line)
|
||||
for line in (root / "marker.jsonl").read_text().splitlines()
|
||||
]
|
||||
for sample in samples:
|
||||
eligible = [
|
||||
marker
|
||||
for marker in markers
|
||||
if marker["monotonic"] <= sample["seconds"]
|
||||
]
|
||||
sample["phase"] = (
|
||||
max(eligible, key=lambda marker: marker["monotonic"])[
|
||||
"phase"
|
||||
]
|
||||
if eligible
|
||||
else "pre-start"
|
||||
)
|
||||
result.update(
|
||||
{
|
||||
"mode": mode,
|
||||
"repeat": repeat,
|
||||
"cold": warmup == 0,
|
||||
"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
|
||||
),
|
||||
"stop_reason": stop_reason,
|
||||
"rss_samples": samples,
|
||||
}
|
||||
)
|
||||
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",
|
||||
"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()
|
||||
if csv_file:
|
||||
csv_file.close()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Measure real image preparation and SDK requests in isolated processes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
from astrbot.core.utils.media_utils import compress_image
|
||||
|
||||
if args.trace_python:
|
||||
tracemalloc.start()
|
||||
process = psutil.Process()
|
||||
baseline_rss = process.memory_info().rss
|
||||
prepared = None
|
||||
phases = {}
|
||||
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:
|
||||
tick = time.perf_counter()
|
||||
prepared = await compress_image(str(args.image.resolve()))
|
||||
phases["prepare_ms"] = (time.perf_counter() - tick) * 1000
|
||||
prepared_bytes = Path(prepared).stat().st_size
|
||||
source_data = args.image.read_bytes()
|
||||
prepared_data = Path(prepared).read_bytes()
|
||||
tick = time.perf_counter()
|
||||
request = ProviderRequest(
|
||||
prompt="Describe this image.", image_urls=[prepared]
|
||||
)
|
||||
context = await request.assemble_context()
|
||||
phases["assemble_base64_ms"] = (time.perf_counter() - tick) * 1000
|
||||
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
|
||||
finally:
|
||||
if prepared and Path(prepared).resolve() != args.image.resolve():
|
||||
Path(prepared).unlink(missing_ok=True)
|
||||
residual = [
|
||||
path
|
||||
for path in Path(runtime).rglob("*")
|
||||
if path.is_file() and path not in before_files
|
||||
]
|
||||
result = {
|
||||
"source_bytes": args.image.stat().st_size,
|
||||
"prepared_bytes": prepared_bytes,
|
||||
"source_sha256": hashlib.sha256(source_data).hexdigest(),
|
||||
"prepared_sha256": hashlib.sha256(prepared_data).hexdigest(),
|
||||
"preserved_source_bytes": prepared_data == source_data,
|
||||
"base64_bytes": 4 * ((prepared_bytes + 2) // 3),
|
||||
"phase_ms": phases,
|
||||
"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"] = "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"}
|
||||
)
|
||||
result.update(
|
||||
{
|
||||
"status": stop_reason or result["status"],
|
||||
"returncode": proc.returncode,
|
||||
"rss_high_water_bytes": max(
|
||||
(s["rss"] for s in samples), 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(),
|
||||
}
|
||||
)
|
||||
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",
|
||||
"post_request_rss",
|
||||
"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"),
|
||||
"post_request_rss": result.get("post_request_rss"),
|
||||
"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())
|
||||
@@ -0,0 +1,99 @@
|
||||
"""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()))
|
||||
@@ -0,0 +1,449 @@
|
||||
#!/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 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 FROM sqlite_master WHERE type IN ('table','index','trigger','view') ORDER BY name"
|
||||
)
|
||||
for name, sql in tables:
|
||||
digest.update(json.dumps([name, sql], ensure_ascii=False).encode())
|
||||
if name.startswith("sqlite_") or sql is None:
|
||||
continue
|
||||
try:
|
||||
rows = connection.execute(
|
||||
f'SELECT * FROM "{name.replace(chr(34), chr(34) * 2)}"'
|
||||
)
|
||||
except sqlite3.DatabaseError:
|
||||
continue
|
||||
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")
|
||||
for part in (
|
||||
item.get("content", [])
|
||||
if isinstance(item, dict)
|
||||
and isinstance(item.get("content"), list)
|
||||
else []
|
||||
):
|
||||
if (
|
||||
isinstance(part, dict)
|
||||
and part.get("type") == "image_media_ref"
|
||||
):
|
||||
ref = ImageMediaStore(media).read
|
||||
# Constructing and reading the reference validates both files and hashes.
|
||||
from astrbot.core.utils.image_media_store import (
|
||||
ImageMediaRef,
|
||||
)
|
||||
|
||||
parsed = ImageMediaRef(
|
||||
**{
|
||||
key: part[key]
|
||||
for key in (
|
||||
"media_id",
|
||||
"mime_type",
|
||||
"width",
|
||||
"height",
|
||||
"byte_size",
|
||||
"detail",
|
||||
"version",
|
||||
"image_id",
|
||||
)
|
||||
}
|
||||
)
|
||||
ref(parsed, {parsed.media_id})
|
||||
referenced.add(parsed.media_id)
|
||||
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())
|
||||
@@ -0,0 +1,59 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,301 @@
|
||||
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_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-*"))
|
||||
Reference in New Issue
Block a user