mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
feat(storage): persist viewed images as shared blobs (#6160)
Store validated image bytes in the configured content-addressed blob backend and resolve checkpointed refs before local compatibility paths. Refs: #4189 Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> Co-authored-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com>
This commit is contained in:
parent
22623d1b00
commit
fcee85b822
@ -4,7 +4,7 @@ Resolves the multi-instance half of [#4189](https://github.com/bytedance/deer-fl
|
||||
|
||||
| Producer | Write | Reads |
|
||||
|---|---|---|
|
||||
| Viewed images | `view_image_tool` → `ViewedImageData.actual_path` (`deerflow/agents/thread_state.py:52`) | `ViewImageMiddleware._read_image_as_data_url`, gateway artifact routes, IM channels, `present_file_tool` |
|
||||
| Viewed images | `view_image_tool` → `ViewedImageData.blob_ref` plus compatibility `actual_path` | `ViewImageMiddleware._read_image_as_data_url`, gateway artifact routes, IM channels, `present_file_tool` |
|
||||
| Externalized tool results | `ToolOutputBudgetMiddleware` → server-owned `tool-output` ref beside the virtual path | pre-model restoration into the current thread-data mount, then model `read_file` |
|
||||
|
||||
On a single gateway both are correct. Behind a load balancer, the instance handling the read is frequently not the instance that wrote the file. The blob store replaces **"where on this machine"** with **"which content"**, so every instance that can reach the backing store resolves the same bytes.
|
||||
@ -59,10 +59,14 @@ key behaves exactly as before.
|
||||
|
||||
The producer migrations are independently revertible:
|
||||
|
||||
1. **Viewed images — pending.** `ViewedImageData` will gain an optional
|
||||
`blob_ref` (`actual_path` is kept); `view_image_tool` will write the blob
|
||||
when the store is enabled; `ViewImageMiddleware._read_image_as_data_url`
|
||||
will resolve blob-first, path-second.
|
||||
1. **Viewed images — migrated.** `ViewedImageData` carries an optional
|
||||
`blob_ref` while retaining `actual_path`; `view_image_tool` writes validated
|
||||
bytes when the store is enabled, and `ViewImageMiddleware` resolves the
|
||||
digest-bound blob before the sandbox/local compatibility paths. The gateway
|
||||
artifact routes and IM channels keep their local-path reads until they can be
|
||||
exercised against a multi-instance deployment. Store initialization or write
|
||||
failures return a generic tool error; read failures fall through to the
|
||||
validated sandbox/host compatibility paths.
|
||||
2. **Externalized tool results — migrated.** Host externalization stores the
|
||||
exact UTF-8 bytes as `kind="tool-output"` and checkpoints a versioned,
|
||||
server-owned ref beside the virtual `read_file` path. Before a model call,
|
||||
|
||||
@ -103,7 +103,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
TODO attribution skips non-string statuses before membership; preserve args/TODO state for normal tool validation.
|
||||
24. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint.
|
||||
25. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket
|
||||
26. **ViewImageMiddleware** - *(vision models only)* Rebuilds hidden base64 image context in `ModelRequest.messages`, never in new checkpoint state; state retains lightweight `viewed_images` metadata. `_image_injection_plan` first sweeps stale server-marked context (including payloads stranded by the old `before_model`/`after_model` lifecycle), then shares one read-eligibility predicate between sync injection and async authorization. The sweep requires both the reserved ID prefix and server-owned marker, preserving user messages and unmarked legacy payloads. Before reading images, both paths recheck `sandbox:execute`, including restored views after role changes; denial skips I/O and removes marked stale context. Host-copy reads additionally require the saved `actual_path` to match the current user/thread's canonical virtual image path, including for old checkpoints and replacement sandboxes; a missing or mismatched thread identity skips that host read. Only the async wrapper passes `authorization_checked=True` after its gate, avoiding a duplicate synchronous provider call in its worker.
|
||||
26. **ViewImageMiddleware** - *(vision models only)* Rebuilds hidden base64 image context in `ModelRequest.messages`, never in new checkpoint state; state retains lightweight `viewed_images` metadata. Blob refs resolve first and fail over to existing paths. `_image_injection_plan` first sweeps stale server-marked context (including payloads stranded by the old `before_model`/`after_model` lifecycle), then shares one read-eligibility predicate between sync injection and async authorization. The sweep requires both the reserved ID prefix and server-owned marker, preserving user messages and unmarked legacy payloads. Before reading images, both paths recheck `sandbox:execute`, including restored views after role changes; denial skips I/O and removes marked stale context. Host-copy reads additionally require the saved `actual_path` to match the current user/thread's canonical virtual image path, including for old checkpoints and replacement sandboxes; a missing or mismatched thread identity skips that host read. Only the async wrapper passes `authorization_checked=True` after its gate, avoiding a duplicate synchronous provider call in its worker.
|
||||
27. **McpRoutingMiddleware** - Auto-promotes deferred schemas matching the latest real user message before `DeferredToolFilterMiddleware`, without executing tools. New names emit `middleware:tool_promotion` with `source=routing_hint`; repeat passes emit nothing
|
||||
28. **DeferredToolPromotionAuditMiddleware** - Observes final `tool_search` `Command`s; keep it outer of `SkillToolPolicyMiddleware` so denied names are excluded. It atomically claims new names per lead run or subagent execution to dedupe parallel searches, derives subagent attribution from the server-installed recorder, returns the original `Command`, and omits private payloads
|
||||
29. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` or `McpRoutingMiddleware` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped)
|
||||
|
||||
@ -167,6 +167,46 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _read_blob_image_as_data_url(
|
||||
cls,
|
||||
blob_ref_data: Mapping[str, object] | None,
|
||||
mime_type: str,
|
||||
expected_size: int,
|
||||
expected_sha256: str | None,
|
||||
) -> str | None:
|
||||
"""Resolve a trusted checkpoint blob ref, rejecting metadata drift."""
|
||||
if blob_ref_data is None or expected_sha256 is None:
|
||||
return None
|
||||
|
||||
from deerflow.storage import BlobRef, get_blob_store_if_enabled
|
||||
|
||||
try:
|
||||
ref = BlobRef.model_validate(dict(blob_ref_data))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if ref.kind != "viewed-image" or ref.sha256 != expected_sha256 or ref.size != expected_size or ref.content_type != mime_type:
|
||||
return None
|
||||
|
||||
try:
|
||||
store = get_blob_store_if_enabled()
|
||||
if store is None:
|
||||
return None
|
||||
image_bytes = store.get_bytes(ref)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to resolve viewed image blob %s",
|
||||
ref.sha256[:12],
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
return cls._encode_image_bytes(
|
||||
image_bytes,
|
||||
mime_type,
|
||||
expected_size,
|
||||
expected_sha256,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _read_image_as_data_url(
|
||||
cls,
|
||||
@ -177,6 +217,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_size: int,
|
||||
expected_sha256: str | None,
|
||||
source_sandbox_id: str | None,
|
||||
blob_ref_data: Mapping[str, object] | None = None,
|
||||
*,
|
||||
allow_host_copy: bool = False,
|
||||
) -> str | None:
|
||||
@ -190,6 +231,15 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
fallback. When no live sandbox exists, the historical host compatibility
|
||||
path remains available (digest-checked when present).
|
||||
"""
|
||||
blob_data_url = cls._read_blob_image_as_data_url(
|
||||
blob_ref_data,
|
||||
mime_type,
|
||||
expected_size,
|
||||
expected_sha256,
|
||||
)
|
||||
if blob_data_url is not None:
|
||||
return blob_data_url
|
||||
|
||||
from deerflow.sandbox.overwrite import unwrap_sandbox
|
||||
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
||||
|
||||
@ -293,6 +343,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_size = image_data.get("size", 0)
|
||||
expected_sha256 = image_data.get("sha256")
|
||||
source_sandbox_id = image_data.get("source_sandbox_id")
|
||||
blob_ref = image_data.get("blob_ref")
|
||||
|
||||
# Add text description
|
||||
content_blocks.append({"type": "text", "text": f"\n- **{image_path}** ({mime_type})"})
|
||||
@ -306,6 +357,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_size,
|
||||
expected_sha256 if isinstance(expected_sha256, str) else None,
|
||||
source_sandbox_id if isinstance(source_sandbox_id, str) else None,
|
||||
blob_ref if isinstance(blob_ref, Mapping) else None,
|
||||
allow_host_copy=host_path_allowed(image_path, actual_path) if host_path_allowed is not None else False,
|
||||
)
|
||||
if data_url:
|
||||
|
||||
@ -50,6 +50,15 @@ class BackgroundTaskState(TypedDict):
|
||||
updated_at: str
|
||||
|
||||
|
||||
class BlobRefData(TypedDict):
|
||||
"""JSON-safe checkpoint representation of ``deerflow.storage.BlobRef``."""
|
||||
|
||||
sha256: str
|
||||
size: int
|
||||
kind: str
|
||||
content_type: NotRequired[str | None]
|
||||
|
||||
|
||||
class ViewedImageData(TypedDict):
|
||||
"""Metadata for a viewed image file.
|
||||
|
||||
@ -64,6 +73,7 @@ class ViewedImageData(TypedDict):
|
||||
actual_path: str
|
||||
sha256: str
|
||||
source_sandbox_id: NotRequired[str]
|
||||
blob_ref: NotRequired[BlobRefData]
|
||||
|
||||
|
||||
def merge_sandbox(existing: SandboxState | None, new: SandboxState | None) -> SandboxState | None:
|
||||
|
||||
@ -21,6 +21,14 @@ Replaced stores stay open for callers that already hold them and are closed by
|
||||
`reset_blob_store()`; changing a root requires preserving existing blob refs
|
||||
and coordinating deployment instances.
|
||||
|
||||
`view_image_tool` is the first migrated producer: it persists validated image
|
||||
bytes as `kind="viewed-image"`, checkpoints the JSON-safe ref beside
|
||||
`actual_path`, and treats any store initialization or write failure as a failed
|
||||
tool call rather than claiming cross-instance durability. `ViewImageMiddleware` accepts a
|
||||
ref only when its kind, digest, size and content type agree with the separately
|
||||
recorded image metadata; reads are blob-first and retain the existing
|
||||
sandbox/host compatibility path as fallback after any store failure.
|
||||
|
||||
`ToolOutputBudgetMiddleware` is a migrated producer for host-externalized
|
||||
results. It persists exact UTF-8 bytes as `kind="tool-output"`, checkpoints a
|
||||
versioned ref beside the virtual path, and restores the file atomically under
|
||||
|
||||
@ -32,7 +32,7 @@ Only standalone tool discovery without a model falls back to the base profile.
|
||||
3. **Built-in tools**:
|
||||
- `present_files` - Make output files visible to user (only `/mnt/user-data/outputs`); virtual paths use `resolve_runtime_user_id(runtime)` so validation resolves the same user-scoped outputs directory established by `ThreadDataMiddleware`
|
||||
- `ask_clarification` - Request clarification (intercepted by ClarificationMiddleware, which preserves text fallback and adds `artifact.human_input` for Web UI Human Input Cards). Beyond free text and single choice, the request-side v2 protocol supports `fields` (structured form card collecting several values at once; field types: text/textarea/number/select/multi_select/checkbox/date, validated and normalized server-side in the middleware — invalid entries are dropped, unknown types degrade to `text`; a standalone multi-select question is a one-field form). Replies stay on the v1 response protocol (`text`/`option`): the form card submits a readable text summary
|
||||
- `view_image` - Read image bytes for vision-capable models; both sync and async entry points require `sandbox:execute` before any host or sandbox read. Live sandbox bytes win for the same sandbox generation, replacement-sandbox recovery uses only SHA-256-verified synchronized host bytes, and async tool invocation drains blocking reads before cancellation may release the sandbox lease
|
||||
- `view_image` - Read image bytes for vision-capable models; both sync and async entry points require `sandbox:execute` before any host or sandbox read. With blob storage enabled, validated bytes are persisted as `kind="viewed-image"` and the checkpoint carries a JSON-safe ref beside the compatibility host path; store initialization/write failures return a generic tool error. Blob-backed model reads validate the ref against the separate image metadata before use and fall through to validated compatibility paths on any store failure. Live sandbox bytes remain the fallback for the same sandbox generation, replacement-sandbox recovery uses only SHA-256-verified synchronized host bytes, and async tool invocation drains blocking reads before cancellation may release the sandbox lease
|
||||
- `setup_agent` - Bootstrap-only: persist a custom agent's `SOUL.md` and `config.yaml`. Re-bootstrapping preserves the owner's existing `display_name`. Bound only when `is_bootstrap=True`.
|
||||
- `update_agent` - Custom-agent-only: persist self-updates to the current agent's `SOUL.md` / `config.yaml` from inside a normal chat (partial update + atomic write). Bound when `agent_name` is set and `is_bootstrap=False`.
|
||||
4. **Subagent tool** (if enabled):
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
import hashlib
|
||||
import logging
|
||||
import mimetypes
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
@ -8,10 +10,12 @@ from langchain_core.messages import ToolMessage
|
||||
from langchain_core.tools import StructuredTool
|
||||
from langgraph.types import Command
|
||||
|
||||
from deerflow.agents.thread_state import ThreadDataState
|
||||
from deerflow.agents.thread_state import ThreadDataState, ViewedImageData
|
||||
from deerflow.config.paths import VIRTUAL_PATH_PREFIX
|
||||
from deerflow.tools.types import Runtime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ALLOWED_IMAGE_VIRTUAL_ROOTS = (
|
||||
f"{VIRTUAL_PATH_PREFIX}/workspace",
|
||||
f"{VIRTUAL_PATH_PREFIX}/uploads",
|
||||
@ -50,6 +54,19 @@ def _sanitize_image_error(error: Exception, thread_data: ThreadDataState | None)
|
||||
return mask_local_paths_in_output(f"{type(error).__name__}: {error}", thread_data)
|
||||
|
||||
|
||||
def _get_thread_id(runtime: Runtime) -> str | None:
|
||||
"""Resolve advisory blob provenance without depending on ambient config."""
|
||||
context = getattr(runtime, "context", None)
|
||||
if isinstance(context, Mapping):
|
||||
thread_id = context.get("thread_id")
|
||||
if isinstance(thread_id, str) and thread_id:
|
||||
return thread_id
|
||||
config = getattr(runtime, "config", None)
|
||||
configurable = config.get("configurable") if isinstance(config, Mapping) else None
|
||||
thread_id = configurable.get("thread_id") if isinstance(configurable, Mapping) else None
|
||||
return thread_id if isinstance(thread_id, str) and thread_id else None
|
||||
|
||||
|
||||
def _is_file_not_found_error(error: BaseException) -> bool:
|
||||
"""Recognize an explicit missing-file signal through provider wrappers.
|
||||
|
||||
@ -274,12 +291,36 @@ def _view_image_authorized(runtime: Runtime, image_path: str, tool_call_id: str)
|
||||
)
|
||||
mime_type = detected_mime_type
|
||||
|
||||
image_metadata = {
|
||||
image_metadata: ViewedImageData = {
|
||||
"mime_type": mime_type,
|
||||
"size": image_size,
|
||||
"actual_path": str(actual_path),
|
||||
"sha256": hashlib.sha256(image_data).hexdigest(),
|
||||
}
|
||||
from deerflow.storage import get_blob_store_if_enabled
|
||||
|
||||
try:
|
||||
blob_store = get_blob_store_if_enabled()
|
||||
if blob_store is not None:
|
||||
blob_ref = blob_store.put_bytes(
|
||||
image_data,
|
||||
kind="viewed-image",
|
||||
content_type=mime_type,
|
||||
thread_id=_get_thread_id(runtime),
|
||||
)
|
||||
image_metadata["blob_ref"] = blob_ref.model_dump(mode="json", exclude_none=True)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist viewed image in shared blob storage", exc_info=True)
|
||||
return Command(
|
||||
update={
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
"Error: Failed to persist image in shared blob storage",
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
]
|
||||
},
|
||||
)
|
||||
if read_source_sandbox_id is not None:
|
||||
image_metadata["source_sandbox_id"] = read_source_sandbox_id
|
||||
new_viewed_images = {image_path: image_metadata}
|
||||
|
||||
@ -22,6 +22,7 @@ Covered behavior:
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
@ -38,6 +39,7 @@ from deerflow.agents.middlewares.view_image_middleware import (
|
||||
ViewImageMiddleware,
|
||||
)
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.storage import BlobRef
|
||||
|
||||
|
||||
def _view_image_call(call_id: str = "call_1", path: str = "/mnt/user-data/uploads/img.png") -> dict:
|
||||
@ -95,6 +97,42 @@ def _authorized_host_view(tmp_path, monkeypatch):
|
||||
return virtual_path, _make_viewed_image(outputs)
|
||||
|
||||
|
||||
def _blob_backed_request(monkeypatch):
|
||||
from deerflow.authz import sandbox_authz
|
||||
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nshared-model-image"
|
||||
ref = BlobRef(
|
||||
sha256=hashlib.sha256(image_bytes).hexdigest(),
|
||||
size=len(image_bytes),
|
||||
kind="viewed-image",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
class SharedStore:
|
||||
def get_bytes(self, requested: BlobRef) -> bytes:
|
||||
assert requested == ref
|
||||
return image_bytes
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: SharedStore())
|
||||
monkeypatch.setattr(sandbox_authz, "authorize_sandbox_execution", lambda **kwargs: None)
|
||||
virtual_path = "/mnt/user-data/outputs/shared.png"
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1", virtual_path)])
|
||||
request = _model_request(
|
||||
[assistant, ToolMessage(content="ok", tool_call_id="c1")],
|
||||
{
|
||||
virtual_path: {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes),
|
||||
"actual_path": "Z:/writer-only/shared.png",
|
||||
"sha256": ref.sha256,
|
||||
"blob_ref": ref.model_dump(exclude_none=True),
|
||||
}
|
||||
},
|
||||
)
|
||||
request.runtime.context = {"user_id": "user-test", "thread_id": "thread-test"}
|
||||
return request, "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
|
||||
|
||||
class TestGetLastAssistantMessage:
|
||||
def test_returns_none_on_empty_list(self):
|
||||
mw = ViewImageMiddleware()
|
||||
@ -236,6 +274,167 @@ class TestCreateImageDetailsMessage:
|
||||
assert blocks[2]["type"] == "image_url"
|
||||
assert blocks[2]["image_url"]["url"].startswith("data:image/png;base64,")
|
||||
|
||||
def test_resolves_blob_ref_when_writer_host_copy_is_absent(self, tmp_path, monkeypatch):
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nshared-image"
|
||||
ref = BlobRef(
|
||||
sha256=hashlib.sha256(image_bytes).hexdigest(),
|
||||
size=len(image_bytes),
|
||||
kind="viewed-image",
|
||||
content_type="image/png",
|
||||
)
|
||||
reads: list[BlobRef] = []
|
||||
|
||||
class SharedStore:
|
||||
def get_bytes(self, requested: BlobRef) -> bytes:
|
||||
reads.append(requested)
|
||||
return image_bytes
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: SharedStore())
|
||||
state = {
|
||||
"viewed_images": {
|
||||
"/mnt/user-data/outputs/shared.png": {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes),
|
||||
"actual_path": str(tmp_path / "writer-only" / "shared.png"),
|
||||
"sha256": ref.sha256,
|
||||
"blob_ref": ref.model_dump(exclude_none=True),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state)
|
||||
|
||||
image_blocks = [block for block in blocks if isinstance(block, dict) and block.get("type") == "image_url"]
|
||||
assert len(image_blocks) == 1
|
||||
assert image_blocks[0]["image_url"]["url"] == "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
assert reads == [ref]
|
||||
|
||||
def test_blob_store_factory_failure_falls_back_to_validated_host_copy(self, tmp_path, monkeypatch):
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nlocal-fallback"
|
||||
img_path = tmp_path / "fallback.png"
|
||||
img_path.write_bytes(image_bytes)
|
||||
ref = BlobRef(
|
||||
sha256=hashlib.sha256(image_bytes).hexdigest(),
|
||||
size=len(image_bytes),
|
||||
kind="viewed-image",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
def fail_to_resolve_store():
|
||||
raise ValueError("misconfigured backend")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", fail_to_resolve_store)
|
||||
state = {
|
||||
"viewed_images": {
|
||||
"/mnt/user-data/outputs/fallback.png": {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes),
|
||||
"actual_path": str(img_path),
|
||||
"sha256": ref.sha256,
|
||||
"blob_ref": ref.model_dump(exclude_none=True),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(
|
||||
state,
|
||||
host_path_allowed=lambda _virtual, actual: actual == str(img_path),
|
||||
)
|
||||
|
||||
image_blocks = [block for block in blocks if isinstance(block, dict) and block.get("type") == "image_url"]
|
||||
assert len(image_blocks) == 1
|
||||
assert image_blocks[0]["image_url"]["url"] == "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
|
||||
def test_raw_blob_read_failure_falls_back_to_validated_host_copy(self, tmp_path, monkeypatch):
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nlocal-fallback"
|
||||
img_path = tmp_path / "fallback.png"
|
||||
img_path.write_bytes(image_bytes)
|
||||
ref = BlobRef(
|
||||
sha256=hashlib.sha256(image_bytes).hexdigest(),
|
||||
size=len(image_bytes),
|
||||
kind="viewed-image",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
class FailingStore:
|
||||
def get_bytes(self, requested: BlobRef) -> bytes:
|
||||
raise RuntimeError("raw SDK failure")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: FailingStore())
|
||||
state = {
|
||||
"viewed_images": {
|
||||
"/mnt/user-data/outputs/fallback.png": {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes),
|
||||
"actual_path": str(img_path),
|
||||
"sha256": ref.sha256,
|
||||
"blob_ref": ref.model_dump(exclude_none=True),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(
|
||||
state,
|
||||
host_path_allowed=lambda _virtual, actual: actual == str(img_path),
|
||||
)
|
||||
|
||||
image_blocks = [block for block in blocks if isinstance(block, dict) and block.get("type") == "image_url"]
|
||||
assert len(image_blocks) == 1
|
||||
assert image_blocks[0]["image_url"]["url"] == "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
|
||||
def test_malformed_blob_ref_falls_back_to_validated_host_copy(self, tmp_path):
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nlocal-fallback"
|
||||
img_path = tmp_path / "fallback.png"
|
||||
img_path.write_bytes(image_bytes)
|
||||
state = {
|
||||
"viewed_images": {
|
||||
"/mnt/user-data/outputs/fallback.png": {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes),
|
||||
"actual_path": str(img_path),
|
||||
"sha256": hashlib.sha256(image_bytes).hexdigest(),
|
||||
"blob_ref": {"sha256": "not-a-digest", "size": len(image_bytes), "kind": "viewed-image"},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(
|
||||
state,
|
||||
host_path_allowed=lambda _virtual, actual: actual == str(img_path),
|
||||
)
|
||||
|
||||
assert any(isinstance(block, dict) and block.get("type") == "image_url" for block in blocks)
|
||||
|
||||
def test_blob_ref_must_match_checkpoint_metadata_before_read(self, tmp_path, monkeypatch):
|
||||
image_bytes = b"\x89PNG\r\n\x1a\nshared-image"
|
||||
ref = BlobRef(
|
||||
sha256=hashlib.sha256(image_bytes).hexdigest(),
|
||||
size=len(image_bytes),
|
||||
kind="viewed-image",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
class StoreMustNotBeRead:
|
||||
def get_bytes(self, requested: BlobRef) -> bytes:
|
||||
raise AssertionError(f"unexpected blob read: {requested}")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: StoreMustNotBeRead())
|
||||
state = {
|
||||
"viewed_images": {
|
||||
"/mnt/user-data/outputs/shared.png": {
|
||||
"mime_type": "image/png",
|
||||
"size": len(image_bytes) + 1,
|
||||
"actual_path": str(tmp_path / "missing.png"),
|
||||
"sha256": ref.sha256,
|
||||
"blob_ref": ref.model_dump(exclude_none=True),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state)
|
||||
|
||||
assert all(not (isinstance(block, dict) and block.get("type") == "image_url") for block in blocks)
|
||||
|
||||
def test_builds_blocks_for_multiple_images(self, tmp_path):
|
||||
mw = ViewImageMiddleware()
|
||||
img1 = _make_viewed_image(tmp_path, "a.png", data=b"\x89PNG\r\n\x1a\nfake-png")
|
||||
@ -497,6 +696,17 @@ class TestInject:
|
||||
|
||||
|
||||
class TestWrapModelCall:
|
||||
def test_blob_backed_image_reaches_sync_model_without_writer_disk(self, monkeypatch):
|
||||
request, expected = _blob_backed_request(monkeypatch)
|
||||
seen: list[ModelRequest] = []
|
||||
|
||||
ViewImageMiddleware().wrap_model_call(
|
||||
request,
|
||||
lambda prepared: seen.append(prepared) or AIMessage(content="ok"),
|
||||
)
|
||||
|
||||
assert _image_urls(seen[0]) == [expected]
|
||||
|
||||
def test_handler_receives_the_image_context_message(self, tmp_path):
|
||||
mw = ViewImageMiddleware()
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
|
||||
@ -543,6 +753,26 @@ class TestWrapModelCall:
|
||||
assert result.content == "I can see the image."
|
||||
assert len(_image_context_messages(seen[0].messages)) == 1
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_blob_backed_image_reaches_async_model_without_writer_disk(self, monkeypatch):
|
||||
from deerflow.authz import sandbox_authz
|
||||
|
||||
request, expected = _blob_backed_request(monkeypatch)
|
||||
seen: list[ModelRequest] = []
|
||||
|
||||
async def authorize(**kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(sandbox_authz, "authorize_sandbox_execution_async", authorize)
|
||||
|
||||
async def handler(prepared: ModelRequest) -> AIMessage:
|
||||
seen.append(prepared)
|
||||
return AIMessage(content="ok")
|
||||
|
||||
await ViewImageMiddleware().awrap_model_call(request, handler)
|
||||
|
||||
assert _image_urls(seen[0]) == [expected]
|
||||
|
||||
|
||||
class TestGraphIntegration:
|
||||
def _graph_and_capture(self):
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
@ -6,6 +7,7 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.storage import BlobRef, BlobWriteError
|
||||
from deerflow.tools.builtins.view_image_tool import view_image_tool
|
||||
|
||||
view_image_module = importlib.import_module("deerflow.tools.builtins.view_image_tool")
|
||||
@ -76,6 +78,122 @@ def test_view_image_reads_virtual_uploads_path(tmp_path: Path) -> None:
|
||||
assert viewed_image["actual_path"] == str(image_path)
|
||||
|
||||
|
||||
def test_view_image_persists_blob_ref_when_enabled(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
thread_data = _make_thread_data(tmp_path)
|
||||
image_path = Path(thread_data["uploads_path"]) / "shared.png"
|
||||
image_path.write_bytes(PNG_BYTES)
|
||||
calls: list[dict] = []
|
||||
|
||||
class RecordingStore:
|
||||
def put_bytes(self, data: bytes, **kwargs) -> BlobRef:
|
||||
calls.append({"data": data, **kwargs})
|
||||
return BlobRef(
|
||||
sha256=hashlib.sha256(data).hexdigest(),
|
||||
size=len(data),
|
||||
kind=kwargs["kind"],
|
||||
content_type=kwargs["content_type"],
|
||||
)
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: RecordingStore())
|
||||
|
||||
result = view_image_tool.func(
|
||||
runtime=_make_runtime(thread_data),
|
||||
image_path="/mnt/user-data/uploads/shared.png",
|
||||
tool_call_id="tc-shared",
|
||||
)
|
||||
|
||||
assert _message_content(result) == "Successfully read image"
|
||||
metadata = result.update["viewed_images"]["/mnt/user-data/uploads/shared.png"]
|
||||
assert metadata["blob_ref"] == {
|
||||
"sha256": hashlib.sha256(PNG_BYTES).hexdigest(),
|
||||
"size": len(PNG_BYTES),
|
||||
"kind": "viewed-image",
|
||||
"content_type": "image/png",
|
||||
}
|
||||
assert calls == [
|
||||
{
|
||||
"data": PNG_BYTES,
|
||||
"kind": "viewed-image",
|
||||
"content_type": "image/png",
|
||||
"thread_id": "thread-1",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_view_image_does_not_claim_success_when_configured_blob_write_fails(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
thread_data = _make_thread_data(tmp_path)
|
||||
image_path = Path(thread_data["uploads_path"]) / "shared.png"
|
||||
image_path.write_bytes(PNG_BYTES)
|
||||
|
||||
class FailingStore:
|
||||
def put_bytes(self, data: bytes, **kwargs) -> BlobRef:
|
||||
raise BlobWriteError("backend detail must not reach the model")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: FailingStore())
|
||||
|
||||
result = view_image_tool.func(
|
||||
runtime=_make_runtime(thread_data),
|
||||
image_path="/mnt/user-data/uploads/shared.png",
|
||||
tool_call_id="tc-shared-failure",
|
||||
)
|
||||
|
||||
assert _message_content(result) == "Error: Failed to persist image in shared blob storage"
|
||||
assert "viewed_images" not in result.update
|
||||
assert "backend detail" not in _message_content(result)
|
||||
|
||||
|
||||
def test_view_image_returns_generic_error_when_blob_store_factory_fails(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
thread_data = _make_thread_data(tmp_path)
|
||||
image_path = Path(thread_data["uploads_path"]) / "shared.png"
|
||||
image_path.write_bytes(PNG_BYTES)
|
||||
|
||||
def fail_to_resolve_store():
|
||||
raise ValueError("misconfigured backend detail must not reach the model")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", fail_to_resolve_store)
|
||||
|
||||
result = view_image_tool.func(
|
||||
runtime=_make_runtime(thread_data),
|
||||
image_path="/mnt/user-data/uploads/shared.png",
|
||||
tool_call_id="tc-shared-factory-failure",
|
||||
)
|
||||
|
||||
assert _message_content(result) == "Error: Failed to persist image in shared blob storage"
|
||||
assert "viewed_images" not in result.update
|
||||
assert "misconfigured backend detail" not in _message_content(result)
|
||||
|
||||
|
||||
def test_view_image_returns_generic_error_when_backend_leaks_raw_exception(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
thread_data = _make_thread_data(tmp_path)
|
||||
image_path = Path(thread_data["uploads_path"]) / "shared.png"
|
||||
image_path.write_bytes(PNG_BYTES)
|
||||
|
||||
class FailingStore:
|
||||
def put_bytes(self, data: bytes, **kwargs) -> BlobRef:
|
||||
raise RuntimeError("raw SDK detail must not reach the model")
|
||||
|
||||
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: FailingStore())
|
||||
|
||||
result = view_image_tool.func(
|
||||
runtime=_make_runtime(thread_data),
|
||||
image_path="/mnt/user-data/uploads/shared.png",
|
||||
tool_call_id="tc-shared-raw-failure",
|
||||
)
|
||||
|
||||
assert _message_content(result) == "Error: Failed to persist image in shared blob storage"
|
||||
assert "viewed_images" not in result.update
|
||||
assert "raw SDK detail" not in _message_content(result)
|
||||
|
||||
|
||||
def test_view_image_reads_gif(tmp_path: Path) -> None:
|
||||
thread_data = _make_thread_data(tmp_path)
|
||||
image_path = Path(thread_data["uploads_path"]) / "animation.gif"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user