mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(sandbox): scope viewed image host reads to the current thread (#5824)
* fix(sandbox): scope viewed image host reads to the current thread * docs(gateway): trim agent guidance to fit CI limit * fix(images): honor thread data paths for saved host copies * test(images): align authorization fixture with thread paths --------- Co-authored-by: YxinMiracle <“939157765@qq.com”>
This commit is contained in:
YxinMiracle
GitHub
parent
c1149ebd0d
commit
7d9025c052
@ -1282,7 +1282,7 @@ Advanced deployments can enable pluggable authorization with `authorization.enab
|
||||
|
||||
Follow-up suggestions also check `model:use` before calling the selected model, including the default model when no name is supplied. A denied model returns HTTP 403 without an LLM call; authorization-provider failures follow the configured `fail_closed` policy.
|
||||
|
||||
For vision-capable agents, `view_image` and the subsequent model-context image read also require `sandbox:execute`. Allowing the tool name alone does not grant image-file access; a role denied sandbox execution cannot reread previously recorded image metadata after its permissions change.
|
||||
For vision-capable agents, `view_image` and the subsequent model-context image read also require `sandbox:execute`. Allowing the tool name alone does not grant image-file access; a role denied sandbox execution cannot reread previously recorded image metadata after its permissions change. External run input and thread-state updates cannot set `viewed_images` or `thread_data`; when a saved image is read from a host copy, its path must resolve to the current user's thread and the recorded virtual image path.
|
||||
|
||||
Advanced deployments can also extend the agent runtime itself by declaring `AgentMiddleware` classes under `extensions.middlewares` in `config.yaml` or `extensions_config.json`. Each entry is a `module.path:ClassName` string (zero-argument constructor) or an object `{class, kwargs}` whose `kwargs` are passed to the constructor. `kwargs` values must be JSON types (object, array, string, number, boolean, or null); YAML dates and timestamps are coerced to ISO strings so they match JSON. DeerFlow loads the same configured list into the lead-agent and subagent pipelines after their built-in runtime middlewares and loop/token guards, but before the terminal-response/safety/clarification tail, so enterprise forks can add domain guardrails, tool-call governance, or observability hooks without patching the built-in middleware builders. Missing packages, invalid classes, broken modules, and constructor errors fail loudly at agent creation. Treat `config.yaml` and `extensions_config.json` as trusted operator-controlled files: middleware paths are code execution, just like custom tool, model, sandbox, guardrail, MCP server, and MCP interceptor declarations. Gateway skill/MCP toggle endpoints preserve this field but do not expose an API write path for `extensions.middlewares`. Separate lead-only/subagent-only middleware lists are not supported yet.
|
||||
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
### Gateway API (`app/gateway/`)
|
||||
|
||||
Reject external run/state writes with `sandbox`, `thread_data`, or `viewed_images`.
|
||||
|
||||
Studio retains sanitized creation metadata.
|
||||
|
||||
Capability Center's `business` adapter validates bundled-provider credentials
|
||||
|
||||
@ -454,7 +454,8 @@ def _normalize_input_messages(
|
||||
def strip_server_owned_state_metadata(values: Mapping[str, Any]) -> dict[str, Any]:
|
||||
"""Validate and sanitize caller-supplied state values before checkpointing.
|
||||
|
||||
The server-owned ``sandbox`` channel is rejected. The ``messages`` channel
|
||||
The server-owned ``sandbox``, ``thread_data``, and ``viewed_images`` channels
|
||||
are rejected. The ``messages`` channel
|
||||
is canonicalized to a list of ``BaseMessage``
|
||||
objects, rejects external system/developer roles with HTTP 400, and strips
|
||||
server-owned metadata. Other channels keep their existing shapes while
|
||||
@ -466,10 +467,12 @@ def strip_server_owned_state_metadata(values: Mapping[str, Any]) -> dict[str, An
|
||||
transform trails, or privileged message roles. Every channel is walked
|
||||
because middleware-contributed channels can also carry message-like values.
|
||||
"""
|
||||
if "sandbox" in values:
|
||||
server_owned_channels = {"sandbox", "thread_data", "viewed_images"}
|
||||
rejected = server_owned_channels.intersection(values)
|
||||
if rejected:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="External sandbox state is not allowed",
|
||||
detail=f"External {sorted(rejected)[0]} state is not allowed",
|
||||
)
|
||||
|
||||
stripped: dict[str, Any] = {}
|
||||
@ -501,9 +504,10 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool
|
||||
of bubbling up as a 500. The gateway is a system boundary, so per-entry
|
||||
validation errors are the right shape for clients to retry against.
|
||||
|
||||
The ``sandbox`` channel is also server-owned. External callers cannot select
|
||||
a provider resource by id; trusted internal run admission may carry the
|
||||
server's own restored value.
|
||||
The ``sandbox``, ``thread_data``, and ``viewed_images`` channels are also
|
||||
server-owned. External callers cannot select a provider resource by id or
|
||||
supply host image paths; trusted internal run admission may carry restored
|
||||
values.
|
||||
|
||||
``original_user_content``, dynamic-context reminder markers, the transient
|
||||
view-image context marker, the execution-only knowledge-scope marker, tool
|
||||
@ -527,11 +531,14 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool
|
||||
"""
|
||||
if raw_input is None:
|
||||
return {}
|
||||
if not trusted_internal and "sandbox" in raw_input:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="External sandbox state is not allowed",
|
||||
)
|
||||
if not trusted_internal:
|
||||
server_owned_channels = {"sandbox", "thread_data", "viewed_images"}
|
||||
rejected = server_owned_channels.intersection(raw_input)
|
||||
if rejected:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"External {sorted(rejected)[0]} state is not allowed",
|
||||
)
|
||||
result = raw_input
|
||||
messages = raw_input.get("messages")
|
||||
if messages is not None:
|
||||
|
||||
@ -109,7 +109,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
21. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable.
|
||||
22. **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.
|
||||
23. **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
|
||||
24. **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. Only the async wrapper passes `authorization_checked=True` after its gate, avoiding a duplicate synchronous provider call in its worker.
|
||||
24. **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.
|
||||
25. **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
|
||||
26. **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
|
||||
27. **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)
|
||||
|
||||
@ -143,11 +143,11 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
) -> str | None:
|
||||
"""Read a validated host mirror and return a data URL, or None on failure.
|
||||
|
||||
``actual_path`` is server-set by ``view_image_tool`` and held in
|
||||
LangGraph-controlled state. The host path remains the compatibility path
|
||||
for local execution and older checkpoints. Provenance-aware checkpoints
|
||||
additionally verify the exact SHA-256 before a synchronized host copy can
|
||||
stand in for bytes from an earlier sandbox generation.
|
||||
Callers must first bind ``actual_path`` to the current run's user,
|
||||
thread, and virtual image path. The host path remains the compatibility
|
||||
path for local execution and older checkpoints. Provenance-aware
|
||||
checkpoints additionally verify the exact SHA-256 before a synchronized
|
||||
host copy can stand in for bytes from an earlier sandbox generation.
|
||||
"""
|
||||
try:
|
||||
file_path = Path(actual_path)
|
||||
@ -177,6 +177,8 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_size: int,
|
||||
expected_sha256: str | None,
|
||||
source_sandbox_id: str | None,
|
||||
*,
|
||||
allow_host_copy: bool = False,
|
||||
) -> str | None:
|
||||
"""Read the exact image bytes represented by ``viewed_images`` metadata.
|
||||
|
||||
@ -204,7 +206,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
# image was originally read from the host). Reproduce the exact
|
||||
# historical bytes rather than letting an unrelated same-path
|
||||
# file in the replacement sandbox win.
|
||||
if actual_path:
|
||||
if actual_path and allow_host_copy:
|
||||
host_data_url = cls._read_host_image_as_data_url(
|
||||
actual_path,
|
||||
mime_type,
|
||||
@ -248,7 +250,7 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_sha256,
|
||||
)
|
||||
|
||||
if not actual_path:
|
||||
if not actual_path or not allow_host_copy:
|
||||
return None
|
||||
return cls._read_host_image_as_data_url(
|
||||
actual_path,
|
||||
@ -257,7 +259,12 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
expected_sha256,
|
||||
)
|
||||
|
||||
def _create_image_details_message(self, state: ViewImageMiddlewareState) -> list[str | dict]:
|
||||
def _create_image_details_message(
|
||||
self,
|
||||
state: ViewImageMiddlewareState,
|
||||
*,
|
||||
host_path_allowed: Callable[[str, str], bool] | None = None,
|
||||
) -> list[str | dict]:
|
||||
"""Create a formatted message with all viewed image details.
|
||||
|
||||
Reads image files on-demand from the active sandbox when available and
|
||||
@ -299,6 +306,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,
|
||||
allow_host_copy=host_path_allowed(image_path, actual_path) if host_path_allowed is not None else False,
|
||||
)
|
||||
if data_url:
|
||||
content_blocks.append(
|
||||
@ -370,6 +378,52 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
context = getattr(request.runtime, "context", None)
|
||||
return context if isinstance(context, Mapping) else {}
|
||||
|
||||
@classmethod
|
||||
def _host_path_matches_request(cls, request: ModelRequest, image_path: str, actual_path: str) -> bool:
|
||||
"""Bind a stored host copy to the authenticated run's user and thread."""
|
||||
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
from deerflow.sandbox.tools import resolve_and_validate_user_data_path, validate_local_tool_path
|
||||
from deerflow.tools.builtins.view_image_tool import _is_allowed_image_virtual_path
|
||||
|
||||
if not isinstance(image_path, str) or not isinstance(actual_path, str) or not _is_allowed_image_virtual_path(image_path):
|
||||
return False
|
||||
|
||||
from langgraph.config import get_config
|
||||
|
||||
context_thread_id = cls._authorization_context(request).get("thread_id")
|
||||
try:
|
||||
configurable = get_config().get("configurable")
|
||||
configured_thread_id = configurable.get("thread_id") if isinstance(configurable, Mapping) else None
|
||||
except RuntimeError:
|
||||
configured_thread_id = None
|
||||
if configured_thread_id and context_thread_id and configured_thread_id != context_thread_id:
|
||||
return False
|
||||
thread_id = configured_thread_id or context_thread_id
|
||||
if not isinstance(thread_id, str) or not thread_id:
|
||||
return False
|
||||
|
||||
try:
|
||||
user_id = resolve_runtime_user_id(request.runtime)
|
||||
thread_data = (request.state or {}).get("thread_data")
|
||||
if isinstance(thread_data, Mapping):
|
||||
# ThreadDataMiddleware may use a custom Paths base. Bind its
|
||||
# selected root to this run before using it to resolve a host copy.
|
||||
root_name = image_path.removeprefix(f"{VIRTUAL_PATH_PREFIX}/").split("/", 1)[0]
|
||||
root_path = thread_data.get(f"{root_name}_path")
|
||||
if not isinstance(root_path, str) or not root_path:
|
||||
return False
|
||||
root = Path(root_path).resolve()
|
||||
if root.parts[-6:] != ("users", user_id, "threads", thread_id, "user-data", root_name):
|
||||
return False
|
||||
validate_local_tool_path(image_path, thread_data, read_only=True)
|
||||
expected_path = Path(resolve_and_validate_user_data_path(image_path, thread_data))
|
||||
else:
|
||||
expected_path = get_paths().resolve_virtual_path(thread_id, image_path, user_id=user_id)
|
||||
return Path(actual_path).resolve() == expected_path
|
||||
except (OSError, PermissionError, TypeError, ValueError):
|
||||
return False
|
||||
|
||||
def _image_injection_plan(self, request: ModelRequest) -> tuple[list[AnyMessage], bool, bool]:
|
||||
"""Share message cleanup and read eligibility across sync/async paths."""
|
||||
messages = [message for message in request.messages if not self._is_image_context_message(message)]
|
||||
@ -421,7 +475,10 @@ class ViewImageMiddleware(AgentMiddleware[ViewImageMiddlewareState]):
|
||||
# Mixed content (text + images) for the model only, so hide it from the
|
||||
# chat UI and IM channels (matches the other middleware-injected context
|
||||
# messages) even though it never leaves this request.
|
||||
image_content = self._create_image_details_message(request.state or {})
|
||||
image_content = self._create_image_details_message(
|
||||
request.state or {},
|
||||
host_path_allowed=lambda image_path, actual_path: self._host_path_matches_request(request, image_path, actual_path),
|
||||
)
|
||||
logger.debug("Injecting image details message with images into the model request")
|
||||
|
||||
return request.override(messages=[*messages, self._create_image_context_message(image_content)])
|
||||
|
||||
@ -211,6 +211,30 @@ def test_normalize_input_none():
|
||||
assert normalize_input(None) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("boundary", ["run", "state"])
|
||||
@pytest.mark.parametrize("channel", ["viewed_images", "thread_data"])
|
||||
def test_external_image_runtime_state_is_rejected(boundary, channel):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.gateway.services import normalize_input, strip_server_owned_state_metadata
|
||||
|
||||
transform = normalize_input if boundary == "run" else strip_server_owned_state_metadata
|
||||
with pytest.raises(HTTPException) as error:
|
||||
transform({channel: {"synthetic": "untrusted"}})
|
||||
|
||||
assert error.value.status_code == 400
|
||||
|
||||
|
||||
def test_trusted_internal_run_preserves_image_runtime_state():
|
||||
from app.gateway.services import normalize_input
|
||||
|
||||
viewed_images = {"/mnt/user-data/outputs/chart.png": {"actual_path": "/synthetic/chart.png"}}
|
||||
thread_data = {"outputs_path": "/synthetic/outputs"}
|
||||
result = normalize_input({"viewed_images": viewed_images, "thread_data": thread_data}, trusted_internal=True)
|
||||
assert result["viewed_images"] == viewed_images
|
||||
assert result["thread_data"] == thread_data
|
||||
|
||||
|
||||
def test_normalize_input_with_messages():
|
||||
from app.gateway.services import normalize_input
|
||||
|
||||
|
||||
@ -21,6 +21,7 @@ Covered behavior:
|
||||
checkpoint retains it, even if the run is interrupted mid-call.
|
||||
"""
|
||||
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
@ -36,6 +37,7 @@ from deerflow.agents.middlewares.view_image_middleware import (
|
||||
_IMAGE_CONTEXT_MESSAGE_MARKER_KEY,
|
||||
ViewImageMiddleware,
|
||||
)
|
||||
from deerflow.config.paths import Paths
|
||||
|
||||
|
||||
def _view_image_call(call_id: str = "call_1", path: str = "/mnt/user-data/uploads/img.png") -> dict:
|
||||
@ -84,6 +86,15 @@ def _make_viewed_image(tmp_path, filename="img.png", mime_type="image/png", data
|
||||
}
|
||||
|
||||
|
||||
def _authorized_host_view(tmp_path, monkeypatch):
|
||||
paths = Paths(tmp_path)
|
||||
monkeypatch.setattr("deerflow.config.paths.get_paths", lambda: paths)
|
||||
outputs = paths.sandbox_outputs_dir("thread-test", user_id="user-test")
|
||||
outputs.mkdir(parents=True)
|
||||
virtual_path = "/mnt/user-data/outputs/img.png"
|
||||
return virtual_path, _make_viewed_image(outputs)
|
||||
|
||||
|
||||
class TestGetLastAssistantMessage:
|
||||
def test_returns_none_on_empty_list(self):
|
||||
mw = ViewImageMiddleware()
|
||||
@ -214,7 +225,7 @@ class TestCreateImageDetailsMessage:
|
||||
"/path/to/cat.png": img_meta,
|
||||
}
|
||||
}
|
||||
blocks = mw._create_image_details_message(state)
|
||||
blocks = mw._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path))
|
||||
|
||||
# header text + per-image description text + per-image image_url block
|
||||
assert len(blocks) == 3
|
||||
@ -235,7 +246,7 @@ class TestCreateImageDetailsMessage:
|
||||
"/b.jpg": img2,
|
||||
}
|
||||
}
|
||||
blocks = mw._create_image_details_message(state)
|
||||
blocks = mw._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path))
|
||||
|
||||
# 1 header + (1 description + 1 image_url) per image = 5 blocks
|
||||
assert len(blocks) == 5
|
||||
@ -269,7 +280,7 @@ class TestCreateImageDetailsMessage:
|
||||
"/mystery.bin": img_meta,
|
||||
}
|
||||
}
|
||||
blocks = mw._create_image_details_message(state)
|
||||
blocks = mw._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path))
|
||||
# The description block should mention unknown
|
||||
description_blocks = [b for b in blocks if b.get("type") == "text" and "/mystery.bin" in b.get("text", "")]
|
||||
assert len(description_blocks) == 1
|
||||
@ -404,11 +415,13 @@ class TestInject:
|
||||
request = _model_request([HumanMessage(content="hi")])
|
||||
assert mw._inject(request) is request
|
||||
|
||||
def test_appends_image_context_message_to_request(self, tmp_path):
|
||||
def test_appends_image_context_message_to_request(self, tmp_path, monkeypatch):
|
||||
mw = ViewImageMiddleware()
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
|
||||
virtual_path, image_metadata = _authorized_host_view(tmp_path, monkeypatch)
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1", virtual_path)])
|
||||
original = [assistant, ToolMessage(content="ok", tool_call_id="c1")]
|
||||
request = _model_request(original, {"/img.png": _make_viewed_image(tmp_path)})
|
||||
request = _model_request(original, {virtual_path: image_metadata})
|
||||
request.runtime.context = {"user_id": "user-test", "thread_id": "thread-test"}
|
||||
|
||||
injected_request = mw._inject(request)
|
||||
|
||||
@ -427,17 +440,19 @@ class TestInject:
|
||||
assert injected.id is not None
|
||||
assert injected.id.startswith("view-image-context:")
|
||||
|
||||
def test_replaces_a_stranded_payload_instead_of_stacking_a_second_one(self, tmp_path):
|
||||
def test_replaces_a_stranded_payload_instead_of_stacking_a_second_one(self, tmp_path, monkeypatch):
|
||||
"""A run that died during the model call can leave the old
|
||||
before_model/after_model pair's message checkpointed. Rebuild it rather
|
||||
than adding a second copy on top."""
|
||||
mw = ViewImageMiddleware()
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1")])
|
||||
virtual_path, image_metadata = _authorized_host_view(tmp_path, monkeypatch)
|
||||
assistant = AIMessage(content="", tool_calls=[_view_image_call("c1", virtual_path)])
|
||||
stranded = ViewImageMiddleware._create_image_context_message([{"type": "text", "text": "stale"}])
|
||||
request = _model_request(
|
||||
[assistant, ToolMessage(content="ok", tool_call_id="c1"), stranded],
|
||||
{"/img.png": _make_viewed_image(tmp_path)},
|
||||
{virtual_path: image_metadata},
|
||||
)
|
||||
request.runtime.context = {"user_id": "user-test", "thread_id": "thread-test"}
|
||||
|
||||
injected = _image_context_messages(mw._inject(request).messages)
|
||||
|
||||
@ -538,19 +553,24 @@ class TestGraphIntegration:
|
||||
)
|
||||
return create_agent(model=model, tools=[], middleware=[ViewImageMiddleware()]), capture
|
||||
|
||||
def _input(self, tmp_path):
|
||||
def _input(self, tmp_path, monkeypatch):
|
||||
virtual_path, image_metadata = _authorized_host_view(tmp_path, monkeypatch)
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(content="", tool_calls=[_view_image_call("c1")]),
|
||||
AIMessage(content="", tool_calls=[_view_image_call("c1", virtual_path)]),
|
||||
ToolMessage(content="ok", tool_call_id="c1"),
|
||||
],
|
||||
"viewed_images": {"/img.png": _make_viewed_image(tmp_path)},
|
||||
"viewed_images": {virtual_path: image_metadata},
|
||||
}
|
||||
|
||||
def test_image_context_reaches_the_model_but_never_the_state(self, tmp_path):
|
||||
def test_image_context_reaches_the_model_but_never_the_state(self, tmp_path, monkeypatch):
|
||||
graph, capture = self._graph_and_capture()
|
||||
|
||||
result = graph.invoke(self._input(tmp_path))
|
||||
result = graph.invoke(
|
||||
self._input(tmp_path, monkeypatch),
|
||||
config={"configurable": {"thread_id": "thread-test"}},
|
||||
context={"user_id": "user-test"},
|
||||
)
|
||||
|
||||
model_image_messages = _image_context_messages(capture.messages)
|
||||
assert len(model_image_messages) == 1
|
||||
@ -560,15 +580,33 @@ class TestGraphIntegration:
|
||||
assert _image_context_messages(result["messages"]) == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_graph_matches_sync_behavior(self, tmp_path):
|
||||
async def test_async_graph_matches_sync_behavior(self, tmp_path, monkeypatch):
|
||||
graph, capture = self._graph_and_capture()
|
||||
|
||||
result = await graph.ainvoke(self._input(tmp_path))
|
||||
result = await graph.ainvoke(
|
||||
self._input(tmp_path, monkeypatch),
|
||||
config={"configurable": {"thread_id": "thread-test"}},
|
||||
context={"user_id": "user-test"},
|
||||
)
|
||||
|
||||
assert len(_image_context_messages(capture.messages)) == 1
|
||||
assert any(block.get("type") == "image_url" for block in _image_context_messages(capture.messages)[0].content)
|
||||
assert _image_context_messages(result["messages"]) == []
|
||||
|
||||
def test_graph_preserves_normalized_client_message_with_reserved_prefix(self, tmp_path):
|
||||
def test_context_thread_id_cannot_override_configured_thread(self, tmp_path, monkeypatch):
|
||||
graph, capture = self._graph_and_capture()
|
||||
|
||||
graph.invoke(
|
||||
self._input(tmp_path, monkeypatch),
|
||||
config={"configurable": {"thread_id": "another-thread"}},
|
||||
context={"user_id": "user-test", "thread_id": "thread-test"},
|
||||
)
|
||||
|
||||
model_image_messages = _image_context_messages(capture.messages)
|
||||
assert len(model_image_messages) == 1
|
||||
assert all(block.get("type") != "image_url" for block in model_image_messages[0].content)
|
||||
|
||||
def test_graph_preserves_normalized_client_message_with_reserved_prefix(self, tmp_path, monkeypatch):
|
||||
from app.gateway.services import normalize_input
|
||||
|
||||
client_id = "view-image-context:client-supplied"
|
||||
@ -591,9 +629,9 @@ class TestGraphIntegration:
|
||||
assert _IMAGE_CONTEXT_MESSAGE_MARKER_KEY not in client_message.additional_kwargs
|
||||
|
||||
graph, capture = self._graph_and_capture()
|
||||
graph_input = self._input(tmp_path)
|
||||
graph_input = self._input(tmp_path, monkeypatch)
|
||||
|
||||
result = graph.invoke({**graph_input, "messages": [client_message, *graph_input["messages"]]})
|
||||
result = graph.invoke({**graph_input, "messages": [client_message, *graph_input["messages"]]}, context={"user_id": "user-test", "thread_id": "thread-test"})
|
||||
|
||||
assert any(message.id == client_id for message in capture.messages)
|
||||
assert any(message.id != client_id and message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is True for message in _image_context_messages(capture.messages))
|
||||
@ -601,3 +639,127 @@ class TestGraphIntegration:
|
||||
assert persisted_client.content == "client-authored message"
|
||||
assert persisted_client.additional_kwargs == {"custom": "keep-me"}
|
||||
assert all(message.additional_kwargs.get(_IMAGE_CONTEXT_MESSAGE_MARKER_KEY) is not True for message in result["messages"] if isinstance(message, HumanMessage))
|
||||
|
||||
|
||||
def test_host_image_copy_is_scoped_to_current_user_and_thread(tmp_path, monkeypatch):
|
||||
"""A stale viewed_images entry must not open another user's host copy."""
|
||||
from deerflow.authz import sandbox_authz
|
||||
|
||||
monkeypatch.setattr(sandbox_authz, "authorize_sandbox_execution", lambda **kwargs: None)
|
||||
paths = Paths(tmp_path)
|
||||
monkeypatch.setattr("deerflow.config.paths.get_paths", lambda: paths)
|
||||
monkeypatch.setattr("deerflow.sandbox.sandbox_provider.get_sandbox_provider", lambda: SimpleNamespace(get=lambda sandbox_id: None))
|
||||
|
||||
image_bytes = base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==")
|
||||
virtual_path = "/mnt/user-data/outputs/canary.png"
|
||||
owner_path = paths.sandbox_outputs_dir("thread-b", user_id="user-b") / "canary.png"
|
||||
owner_path.parent.mkdir(parents=True)
|
||||
owner_path.write_bytes(image_bytes)
|
||||
metadata = {"mime_type": "image/png", "size": len(image_bytes), "actual_path": str(owner_path)}
|
||||
messages = [AIMessage(content="", tool_calls=[_view_image_call("call-image", virtual_path)]), ToolMessage(content="Successfully read image", tool_call_id="call-image")]
|
||||
|
||||
def model_payloads(user_id: str, thread_id: str | None, actual_path: Path = owner_path) -> list[str]:
|
||||
request = _model_request(messages, {virtual_path: {**metadata, "actual_path": str(actual_path)}})
|
||||
request.runtime.context = {"user_id": user_id, "thread_id": thread_id}
|
||||
prepared = ViewImageMiddleware()._inject(request)
|
||||
return [block["image_url"]["url"] for message in _image_context_messages(prepared.messages) for block in message.content if isinstance(block, dict) and block.get("type") == "image_url"]
|
||||
|
||||
expected = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
assert model_payloads("user-b", "thread-b") == [expected]
|
||||
assert model_payloads("user-a", "thread-a") == []
|
||||
assert model_payloads("user-b", "thread-a") == []
|
||||
assert model_payloads("user-b", None) == []
|
||||
|
||||
alias_path = paths.sandbox_outputs_dir("thread-a", user_id="user-a") / "canary.png"
|
||||
alias_path.parent.mkdir(parents=True)
|
||||
alias_path.symlink_to(owner_path)
|
||||
assert model_payloads("user-a", "thread-a", alias_path) == []
|
||||
|
||||
|
||||
def _custom_base_view_request(tmp_path, monkeypatch):
|
||||
"""Use the real tool to record a host copy under a non-global Paths base."""
|
||||
from deerflow.authz import sandbox_authz
|
||||
from deerflow.tools.builtins.view_image_tool import _view_image_authorized
|
||||
|
||||
monkeypatch.setattr(sandbox_authz, "authorize_sandbox_execution", lambda **kwargs: None)
|
||||
monkeypatch.setattr("deerflow.sandbox.sandbox_provider.get_sandbox_provider", lambda: SimpleNamespace(get=lambda sandbox_id: None))
|
||||
|
||||
paths = Paths(tmp_path / "custom-base")
|
||||
virtual_path = "/mnt/user-data/outputs/canary.png"
|
||||
image_bytes = base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==")
|
||||
actual_path = paths.sandbox_outputs_dir("thread-test", user_id="user-test") / "canary.png"
|
||||
actual_path.parent.mkdir(parents=True)
|
||||
actual_path.write_bytes(image_bytes)
|
||||
thread_data = {
|
||||
"workspace_path": str(paths.sandbox_work_dir("thread-test", user_id="user-test")),
|
||||
"uploads_path": str(paths.sandbox_uploads_dir("thread-test", user_id="user-test")),
|
||||
"outputs_path": str(actual_path.parent),
|
||||
}
|
||||
tool_runtime = SimpleNamespace(
|
||||
state={"thread_data": thread_data},
|
||||
context={"user_id": "user-test", "thread_id": "thread-test"},
|
||||
)
|
||||
result = _view_image_authorized(tool_runtime, virtual_path, "call-image")
|
||||
metadata = result.update["viewed_images"][virtual_path]
|
||||
assert metadata["actual_path"] == str(actual_path)
|
||||
|
||||
messages = [
|
||||
AIMessage(content="", tool_calls=[_view_image_call("call-image", virtual_path)]),
|
||||
ToolMessage(content="Successfully read image", tool_call_id="call-image"),
|
||||
]
|
||||
request = _model_request(messages, {virtual_path: metadata})
|
||||
request.state["thread_data"] = thread_data
|
||||
request.runtime.context = tool_runtime.context
|
||||
expected = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
||||
return request, expected
|
||||
|
||||
|
||||
def _image_urls(request: ModelRequest) -> list[str]:
|
||||
return [block["image_url"]["url"] for message in _image_context_messages(request.messages) for block in message.content if isinstance(block, dict) and block.get("type") == "image_url"]
|
||||
|
||||
|
||||
def test_custom_base_host_copy_reaches_sync_model(tmp_path, monkeypatch):
|
||||
request, expected = _custom_base_view_request(tmp_path, monkeypatch)
|
||||
prepared: list[ModelRequest] = []
|
||||
|
||||
ViewImageMiddleware().wrap_model_call(request, lambda value: prepared.append(value) or AIMessage(content="ok"))
|
||||
|
||||
assert _image_urls(prepared[0]) == [expected]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_custom_base_host_copy_reaches_async_model(tmp_path, monkeypatch):
|
||||
request, expected = _custom_base_view_request(tmp_path, monkeypatch)
|
||||
prepared: list[ModelRequest] = []
|
||||
|
||||
async def handler(value: ModelRequest) -> AIMessage:
|
||||
prepared.append(value)
|
||||
return AIMessage(content="ok")
|
||||
|
||||
await ViewImageMiddleware().awrap_model_call(request, handler)
|
||||
|
||||
assert _image_urls(prepared[0]) == [expected]
|
||||
|
||||
|
||||
def test_custom_base_host_copy_still_requires_matching_user_and_thread(tmp_path, monkeypatch):
|
||||
request, _ = _custom_base_view_request(tmp_path, monkeypatch)
|
||||
middleware = ViewImageMiddleware()
|
||||
|
||||
for user_id, thread_id in (("other-user", "thread-test"), ("user-test", "other-thread"), ("user-test", None)):
|
||||
request.runtime.context = {"user_id": user_id, "thread_id": thread_id}
|
||||
assert _image_urls(middleware._inject(request)) == []
|
||||
|
||||
|
||||
def test_custom_base_host_copy_rejects_unscoped_thread_data_root(tmp_path, monkeypatch):
|
||||
request, _ = _custom_base_view_request(tmp_path, monkeypatch)
|
||||
virtual_path = "/mnt/user-data/outputs/canary.png"
|
||||
unscoped_path = tmp_path / "unscoped" / "canary.png"
|
||||
unscoped_path.parent.mkdir()
|
||||
unscoped_path.write_bytes(base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="))
|
||||
request.state["thread_data"] = {**request.state["thread_data"], "outputs_path": str(unscoped_path.parent)}
|
||||
request.state["viewed_images"][virtual_path] = {
|
||||
**request.state["viewed_images"][virtual_path],
|
||||
"actual_path": str(unscoped_path),
|
||||
}
|
||||
|
||||
assert _image_urls(ViewImageMiddleware()._inject(request)) == []
|
||||
|
||||
@ -486,7 +486,7 @@ def test_middleware_falls_back_to_host_when_saved_sandbox_has_no_live_client(tmp
|
||||
},
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state)
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path))
|
||||
|
||||
assert _image_bytes_from_blocks(blocks) == PNG_BYTES
|
||||
assert provider.acquire_calls == []
|
||||
@ -512,7 +512,7 @@ def test_middleware_uses_verified_host_copy_after_sandbox_replacement(tmp_path,
|
||||
},
|
||||
}
|
||||
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state)
|
||||
blocks = ViewImageMiddleware()._create_image_details_message(state, host_path_allowed=lambda _, actual: Path(actual).is_relative_to(tmp_path))
|
||||
|
||||
assert _image_bytes_from_blocks(blocks) == PNG_BYTES
|
||||
assert sandbox.downloads == []
|
||||
|
||||
@ -31,6 +31,7 @@ from deerflow.authz.tool_filter import apply_tool_authorization
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.config.authorization_config import AuthorizationConfig
|
||||
from deerflow.config.model_config import ModelConfig
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.config.sandbox_config import SandboxConfig
|
||||
from deerflow.guardrails.provider import GuardrailRequest
|
||||
from deerflow.sandbox.exceptions import SandboxAuthorizationError
|
||||
@ -41,13 +42,21 @@ IMAGE_PATH = "/mnt/user-data/uploads/synthetic.gif"
|
||||
|
||||
|
||||
def _setup(root: Path, monkeypatch: pytest.MonkeyPatch, *, sandbox_allowed: bool):
|
||||
uploads = root / "uploads"
|
||||
uploads.mkdir()
|
||||
context = {"thread_id": "synthetic-thread", "user_id": "synthetic-user", "user_role": "reviewer"}
|
||||
paths = Paths(root)
|
||||
thread_id, user_id = context["thread_id"], context["user_id"]
|
||||
uploads = paths.sandbox_uploads_dir(thread_id, user_id=user_id)
|
||||
uploads.mkdir(parents=True)
|
||||
image_file = uploads / "synthetic.gif"
|
||||
image_file.write_bytes(GIF_BYTES)
|
||||
context = {"thread_id": "synthetic-thread", "user_id": "synthetic-user", "user_role": "reviewer"}
|
||||
runtime = SimpleNamespace(
|
||||
state={"thread_data": {"workspace_path": str(root / "workspace"), "uploads_path": str(uploads), "outputs_path": str(root / "outputs")}},
|
||||
state={
|
||||
"thread_data": {
|
||||
"workspace_path": str(paths.sandbox_work_dir(thread_id, user_id=user_id)),
|
||||
"uploads_path": str(uploads),
|
||||
"outputs_path": str(paths.sandbox_outputs_dir(thread_id, user_id=user_id)),
|
||||
}
|
||||
},
|
||||
context=context,
|
||||
config={},
|
||||
)
|
||||
@ -81,7 +90,7 @@ def _model_request(runtime, viewed_images: dict, tool_message: ToolMessage | Non
|
||||
tool_choice=None,
|
||||
tools=[],
|
||||
response_format=None,
|
||||
state={"messages": messages, "viewed_images": viewed_images},
|
||||
state={"messages": messages, "viewed_images": viewed_images, "thread_data": runtime.state["thread_data"]},
|
||||
runtime=SimpleNamespace(context=runtime.context),
|
||||
model_settings={},
|
||||
)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user