chore: import upstream snapshot with attribution
Validate YAML Workflows / Validate YAML Configuration Files (push) Waiting to run
Validate YAML Workflows / Validate YAML Configuration Files (push) Waiting to run
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
"""Regression tests for issue #638.
|
||||
|
||||
``AttachmentService.save_upload_file`` joined the client-supplied multipart
|
||||
filename onto a freshly created temp directory without sanitization. A filename
|
||||
containing ``../`` segments escaped that directory, so the upload wrote
|
||||
attacker-controlled bytes to an arbitrary host path and the cleanup ``unlink``
|
||||
then deleted that same traversed path -- arbitrary file write and delete behind
|
||||
an unauthenticated endpoint.
|
||||
|
||||
These tests pin the basename normalisation and prove a traversal filename can no
|
||||
longer reach a victim file outside the temp/WareHouse area.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
|
||||
import pytest
|
||||
from fastapi import UploadFile
|
||||
|
||||
from server.services.attachment_service import AttachmentService
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw, expected",
|
||||
[
|
||||
("report.pdf", "report.pdf"),
|
||||
("../../../etc/passwd", "passwd"),
|
||||
("..\\..\\windows\\system32\\drivers\\etc\\hosts", "hosts"),
|
||||
("/abs/path/secret.key", "secret.key"),
|
||||
("nested/dir/photo.png", "photo.png"),
|
||||
("", "upload.bin"),
|
||||
(None, "upload.bin"),
|
||||
("..", "upload.bin"),
|
||||
(".", "upload.bin"),
|
||||
(" ", "upload.bin"),
|
||||
],
|
||||
)
|
||||
def test_safe_upload_filename_strips_directory_components(raw, expected):
|
||||
assert AttachmentService._safe_upload_filename(raw) == expected
|
||||
|
||||
|
||||
def _make_upload(filename: str, data: bytes = b"payload") -> UploadFile:
|
||||
return UploadFile(filename=filename, file=io.BytesIO(data))
|
||||
|
||||
|
||||
def test_traversal_filename_cannot_touch_file_outside_temp_dir(tmp_path):
|
||||
"""A traversal filename must neither overwrite nor delete a victim file."""
|
||||
service = AttachmentService(root=tmp_path / "WareHouse")
|
||||
|
||||
victim = tmp_path / "victim.txt"
|
||||
victim.write_text("do-not-touch")
|
||||
|
||||
# Enough parent segments to climb to the filesystem root from any mkdtemp
|
||||
# location, then descend back to the absolute victim path. Pre-fix this
|
||||
# resolved onto the victim and the cleanup unlinked it.
|
||||
traversal = "../" * 16 + str(victim).lstrip("/")
|
||||
record = asyncio.run(service.save_upload_file("sess1", _make_upload(traversal)))
|
||||
|
||||
# Victim survived untouched.
|
||||
assert victim.exists()
|
||||
assert victim.read_text() == "do-not-touch"
|
||||
|
||||
# The stored attachment used the sanitized basename, not the traversal path.
|
||||
assert record.ref.name == "victim.txt"
|
||||
assert "victim.txt" in record.ref.local_path
|
||||
|
||||
|
||||
def test_normal_upload_still_round_trips(tmp_path):
|
||||
service = AttachmentService(root=tmp_path / "WareHouse")
|
||||
record = asyncio.run(
|
||||
service.save_upload_file("sess2", _make_upload("notes.txt", b"hello world"))
|
||||
)
|
||||
assert record.ref.name == "notes.txt"
|
||||
assert record.ref.size == len(b"hello world")
|
||||
@@ -0,0 +1,451 @@
|
||||
"""Tests for Mem0 memory store implementation."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
|
||||
from entity.configs.node.memory import Mem0MemoryConfig
|
||||
from runtime.node.agent.memory.memory_base import (
|
||||
MemoryContentSnapshot,
|
||||
MemoryItem,
|
||||
MemoryWritePayload,
|
||||
)
|
||||
|
||||
|
||||
def _make_store(user_id=None, agent_id=None, api_key="test-key"):
|
||||
"""Build a minimal MemoryStoreConfig mock for Mem0Memory."""
|
||||
mem0_cfg = MagicMock(spec=Mem0MemoryConfig)
|
||||
mem0_cfg.api_key = api_key
|
||||
mem0_cfg.org_id = None
|
||||
mem0_cfg.project_id = None
|
||||
mem0_cfg.user_id = user_id
|
||||
mem0_cfg.agent_id = agent_id
|
||||
|
||||
store = MagicMock()
|
||||
store.name = "test_mem0"
|
||||
|
||||
# Return correct config type based on the requested class
|
||||
def _as_config_side_effect(expected_type, **kwargs):
|
||||
if expected_type is Mem0MemoryConfig:
|
||||
return mem0_cfg
|
||||
return None
|
||||
|
||||
store.as_config.side_effect = _as_config_side_effect
|
||||
return store
|
||||
|
||||
|
||||
def _make_mem0_memory(user_id=None, agent_id=None):
|
||||
"""Create a Mem0Memory with a mocked client."""
|
||||
with patch("runtime.node.agent.memory.mem0_memory._get_mem0_client") as mock_get:
|
||||
mock_client = MagicMock()
|
||||
mock_get.return_value = mock_client
|
||||
from runtime.node.agent.memory.mem0_memory import Mem0Memory
|
||||
store = _make_store(user_id=user_id, agent_id=agent_id)
|
||||
memory = Mem0Memory(store)
|
||||
return memory, mock_client
|
||||
|
||||
|
||||
class TestMem0MemoryRetrieve:
|
||||
|
||||
def test_retrieve_with_agent_id(self):
|
||||
"""Retrieve passes agent_id in filters dict to SDK search."""
|
||||
memory, client = _make_mem0_memory(agent_id="agent-1")
|
||||
client.search.return_value = {
|
||||
"memories": [
|
||||
{"id": "m1", "memory": "test fact", "score": 0.95},
|
||||
]
|
||||
}
|
||||
|
||||
query = MemoryContentSnapshot(text="what do you know?")
|
||||
results = memory.retrieve("writer", query, top_k=5, similarity_threshold=-1.0)
|
||||
|
||||
client.search.assert_called_once()
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["filters"] == {"agent_id": "agent-1"}
|
||||
assert len(results) == 1
|
||||
assert results[0].content_summary == "test fact"
|
||||
assert results[0].metadata["source"] == "mem0"
|
||||
|
||||
def test_retrieve_with_user_id(self):
|
||||
"""Retrieve passes user_id in filters dict to SDK search."""
|
||||
memory, client = _make_mem0_memory(user_id="user-1")
|
||||
client.search.return_value = {
|
||||
"memories": [
|
||||
{"id": "m1", "memory": "user pref", "score": 0.9},
|
||||
]
|
||||
}
|
||||
|
||||
query = MemoryContentSnapshot(text="preferences")
|
||||
results = memory.retrieve("assistant", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["filters"] == {"user_id": "user-1"}
|
||||
assert len(results) == 1
|
||||
|
||||
def test_retrieve_with_both_ids_uses_or_filter(self):
|
||||
"""When both user_id and agent_id are set, an OR filter is used."""
|
||||
memory, client = _make_mem0_memory(user_id="user-1", agent_id="agent-1")
|
||||
client.search.return_value = {
|
||||
"memories": [
|
||||
{"id": "u1", "memory": "user fact", "score": 0.8},
|
||||
{"id": "a1", "memory": "agent fact", "score": 0.9},
|
||||
]
|
||||
}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
results = memory.retrieve("writer", query, top_k=5, similarity_threshold=-1.0)
|
||||
|
||||
client.search.assert_called_once()
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["filters"] == {
|
||||
"OR": [
|
||||
{"user_id": "user-1"},
|
||||
{"agent_id": "agent-1"},
|
||||
]
|
||||
}
|
||||
assert len(results) == 2
|
||||
|
||||
def test_retrieve_fallback_uses_agent_role(self):
|
||||
"""When no IDs configured, fall back to agent_role as agent_id in filters."""
|
||||
memory, client = _make_mem0_memory()
|
||||
client.search.return_value = {"memories": []}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
memory.retrieve("coder", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["filters"] == {"agent_id": "coder"}
|
||||
|
||||
def test_retrieve_empty_query_returns_empty(self):
|
||||
"""Empty query text returns empty without calling API."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
|
||||
query = MemoryContentSnapshot(text=" ")
|
||||
results = memory.retrieve("writer", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
assert results == []
|
||||
client.search.assert_not_called()
|
||||
|
||||
def test_retrieve_api_error_returns_empty(self):
|
||||
"""API errors are caught and return empty list."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.side_effect = Exception("API down")
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
results = memory.retrieve("writer", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_retrieve_respects_top_k(self):
|
||||
"""top_k is passed to Mem0 search."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.return_value = {"memories": []}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
memory.retrieve("writer", query, top_k=7, similarity_threshold=-1.0)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["top_k"] == 7
|
||||
|
||||
def test_retrieve_passes_threshold_when_non_negative(self):
|
||||
"""Non-negative similarity_threshold is forwarded to Mem0."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.return_value = {"memories": []}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
memory.retrieve("writer", query, top_k=3, similarity_threshold=0.5)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["threshold"] == 0.5
|
||||
|
||||
def test_retrieve_passes_zero_threshold(self):
|
||||
"""A threshold of 0.0 is a valid value and should be sent."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.return_value = {"memories": []}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
memory.retrieve("writer", query, top_k=3, similarity_threshold=0.0)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert call_kwargs["threshold"] == 0.0
|
||||
|
||||
def test_retrieve_skips_threshold_when_negative(self):
|
||||
"""Negative similarity_threshold is not sent to Mem0."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.return_value = {"memories": []}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
memory.retrieve("writer", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
call_kwargs = client.search.call_args[1]
|
||||
assert "threshold" not in call_kwargs
|
||||
|
||||
def test_retrieve_handles_legacy_results_key(self):
|
||||
"""Handles SDK response with 'results' key (older SDK versions)."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.search.return_value = {
|
||||
"results": [
|
||||
{"id": "m1", "memory": "legacy format", "score": 0.8},
|
||||
]
|
||||
}
|
||||
|
||||
query = MemoryContentSnapshot(text="test")
|
||||
results = memory.retrieve("writer", query, top_k=3, similarity_threshold=-1.0)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].content_summary == "legacy format"
|
||||
|
||||
|
||||
class TestMem0MemoryUpdate:
|
||||
|
||||
def test_update_sends_only_user_input(self):
|
||||
"""Update sends only user input, not assistant output, to prevent noise."""
|
||||
memory, client = _make_mem0_memory(agent_id="agent-1")
|
||||
client.add.return_value = [{"id": "new", "event": "ADD"}]
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="Write about AI",
|
||||
input_snapshot=MemoryContentSnapshot(text="Write about AI"),
|
||||
output_snapshot=MemoryContentSnapshot(text="AI is transformative..."),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
client.add.assert_called_once()
|
||||
call_kwargs = client.add.call_args[1]
|
||||
assert call_kwargs["agent_id"] == "agent-1"
|
||||
assert "user_id" not in call_kwargs
|
||||
messages = call_kwargs["messages"]
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "user"
|
||||
assert messages[0]["content"] == "Write about AI"
|
||||
|
||||
def test_update_does_not_send_async_mode(self):
|
||||
"""Update does not send deprecated async_mode parameter."""
|
||||
memory, client = _make_mem0_memory(agent_id="agent-1")
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="test",
|
||||
input_snapshot=None,
|
||||
output_snapshot=MemoryContentSnapshot(text="output"),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
assert "async_mode" not in call_kwargs
|
||||
assert call_kwargs["infer"] is True
|
||||
|
||||
def test_update_with_user_id(self):
|
||||
"""User-scoped update uses user_id, not agent_id."""
|
||||
memory, client = _make_mem0_memory(user_id="user-1")
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="I prefer Python",
|
||||
input_snapshot=None,
|
||||
output_snapshot=None,
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
assert call_kwargs["user_id"] == "user-1"
|
||||
assert "agent_id" not in call_kwargs
|
||||
|
||||
def test_update_fallback_uses_agent_role(self):
|
||||
"""When no IDs configured, uses agent_role as agent_id."""
|
||||
memory, client = _make_mem0_memory()
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="coder",
|
||||
inputs_text="test input",
|
||||
input_snapshot=None,
|
||||
output_snapshot=None,
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
assert call_kwargs["agent_id"] == "coder"
|
||||
|
||||
def test_update_with_both_ids_includes_both(self):
|
||||
"""When both user_id and agent_id configured, both are included in add() call."""
|
||||
memory, client = _make_mem0_memory(user_id="user-1", agent_id="agent-1")
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="input",
|
||||
input_snapshot=None,
|
||||
output_snapshot=None,
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
assert call_kwargs["agent_id"] == "agent-1"
|
||||
assert call_kwargs["user_id"] == "user-1"
|
||||
|
||||
def test_update_empty_input_is_noop(self):
|
||||
"""Empty inputs_text skips API call."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text=" ",
|
||||
input_snapshot=None,
|
||||
output_snapshot=MemoryContentSnapshot(text="some output"),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
client.add.assert_not_called()
|
||||
|
||||
def test_update_no_input_is_noop(self):
|
||||
"""No inputs_text skips API call."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="",
|
||||
input_snapshot=None,
|
||||
output_snapshot=MemoryContentSnapshot(text="output"),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
client.add.assert_not_called()
|
||||
|
||||
def test_update_api_error_does_not_raise(self):
|
||||
"""API errors are logged but do not propagate."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.add.side_effect = Exception("API error")
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="test user input",
|
||||
input_snapshot=None,
|
||||
output_snapshot=None,
|
||||
)
|
||||
# Should not raise
|
||||
memory.update(payload)
|
||||
|
||||
|
||||
class TestMem0MemoryPipelineTextCleaning:
|
||||
|
||||
def test_strips_input_from_task_header(self):
|
||||
"""Pipeline headers like '=== INPUT FROM TASK (user) ===' are stripped."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text="=== INPUT FROM TASK (user) ===\n\nMy name is Alex, I love Python",
|
||||
input_snapshot=None,
|
||||
output_snapshot=MemoryContentSnapshot(text="Nice to meet you Alex!"),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
messages = call_kwargs["messages"]
|
||||
assert messages[0]["role"] == "user"
|
||||
assert messages[0]["content"] == "My name is Alex, I love Python"
|
||||
assert "INPUT FROM" not in messages[0]["content"]
|
||||
|
||||
def test_strips_multiple_input_headers(self):
|
||||
"""Multiple pipeline headers from different sources are all stripped."""
|
||||
memory, client = _make_mem0_memory(agent_id="a1")
|
||||
client.add.return_value = []
|
||||
|
||||
payload = MemoryWritePayload(
|
||||
agent_role="writer",
|
||||
inputs_text=(
|
||||
"=== INPUT FROM TASK (user) ===\n\nHello\n\n"
|
||||
"=== INPUT FROM reviewer (assistant) ===\n\nWorld"
|
||||
),
|
||||
input_snapshot=None,
|
||||
output_snapshot=MemoryContentSnapshot(text="Hi!"),
|
||||
)
|
||||
memory.update(payload)
|
||||
|
||||
call_kwargs = client.add.call_args[1]
|
||||
user_content = call_kwargs["messages"][0]["content"]
|
||||
assert "INPUT FROM" not in user_content
|
||||
assert "Hello" in user_content
|
||||
assert "World" in user_content
|
||||
|
||||
def test_clean_text_without_headers_unchanged(self):
|
||||
"""Text without pipeline headers passes through unchanged."""
|
||||
from runtime.node.agent.memory.mem0_memory import Mem0Memory
|
||||
assert Mem0Memory._clean_pipeline_text("Just normal text") == "Just normal text"
|
||||
|
||||
|
||||
class TestMem0MemoryLoadSave:
|
||||
|
||||
def test_load_is_noop(self):
|
||||
"""load() does nothing for cloud-managed store."""
|
||||
memory, _ = _make_mem0_memory(agent_id="a1")
|
||||
memory.load() # Should not raise
|
||||
|
||||
def test_save_is_noop(self):
|
||||
"""save() does nothing for cloud-managed store."""
|
||||
memory, _ = _make_mem0_memory(agent_id="a1")
|
||||
memory.save() # Should not raise
|
||||
|
||||
|
||||
class TestMem0MemoryConfig:
|
||||
|
||||
def test_config_from_dict(self):
|
||||
"""Config parses from dict correctly."""
|
||||
data = {
|
||||
"api_key": "test-key",
|
||||
"user_id": "u1",
|
||||
"org_id": "org-1",
|
||||
}
|
||||
config = Mem0MemoryConfig.from_dict(data, path="test")
|
||||
assert config.api_key == "test-key"
|
||||
assert config.user_id == "u1"
|
||||
assert config.org_id == "org-1"
|
||||
assert config.agent_id is None
|
||||
assert config.project_id is None
|
||||
|
||||
def test_config_field_specs_exist(self):
|
||||
"""FIELD_SPECS are defined for UI generation."""
|
||||
specs = Mem0MemoryConfig.field_specs()
|
||||
assert "api_key" in specs
|
||||
assert "user_id" in specs
|
||||
assert "agent_id" in specs
|
||||
assert specs["api_key"].required is True
|
||||
|
||||
def test_config_requires_api_key(self):
|
||||
"""Config raises ConfigError when api_key is missing."""
|
||||
from entity.configs.base import ConfigError
|
||||
|
||||
data = {"agent_id": "a1"}
|
||||
with pytest.raises(ConfigError):
|
||||
Mem0MemoryConfig.from_dict(data, path="test")
|
||||
|
||||
|
||||
class TestMem0MemoryConstructor:
|
||||
|
||||
def test_raises_on_wrong_config_type(self):
|
||||
"""Mem0Memory raises ValueError when store has wrong config type."""
|
||||
from runtime.node.agent.memory.mem0_memory import Mem0Memory
|
||||
|
||||
store = MagicMock()
|
||||
store.name = "bad_store"
|
||||
store.as_config.return_value = None # Wrong config type
|
||||
|
||||
with pytest.raises(ValueError, match="Mem0 memory store configuration"):
|
||||
Mem0Memory(store)
|
||||
|
||||
def test_import_error_when_mem0ai_missing(self):
|
||||
"""Helpful ImportError when mem0ai is not installed."""
|
||||
from runtime.node.agent.memory.mem0_memory import _get_mem0_client
|
||||
|
||||
mem0_cfg = MagicMock(spec=Mem0MemoryConfig)
|
||||
mem0_cfg.api_key = "test"
|
||||
mem0_cfg.org_id = None
|
||||
mem0_cfg.project_id = None
|
||||
|
||||
with patch.dict("sys.modules", {"mem0": None}):
|
||||
with pytest.raises(ImportError, match="pip install mem0ai"):
|
||||
_get_mem0_client(mem0_cfg)
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Tests for memory embedding dimension consistency."""
|
||||
from unittest.mock import MagicMock, patch
|
||||
from runtime.node.agent.memory.memory_base import MemoryContentSnapshot, MemoryItem
|
||||
from runtime.node.agent.memory.simple_memory import SimpleMemory
|
||||
|
||||
def _make_store(memory_path=None):
|
||||
"""Build a minimal MemoryStoreConfig mock for SimpleMemory."""
|
||||
simple_cfg = MagicMock()
|
||||
simple_cfg.memory_path = memory_path
|
||||
simple_cfg.embedding = None # We'll set embedding manually
|
||||
|
||||
store = MagicMock()
|
||||
store.name = "test_store"
|
||||
store.as_config.return_value = simple_cfg
|
||||
return store
|
||||
|
||||
|
||||
def _make_embedding(dim: int):
|
||||
"""Create a mock EmbeddingBase that produces vectors of the given dimension."""
|
||||
emb = MagicMock()
|
||||
emb.get_embedding.return_value = [0.1] * dim
|
||||
return emb
|
||||
|
||||
|
||||
def _make_memory_item(item_id: str, dim: int):
|
||||
"""Create a MemoryItem with an embedding of the specified dimension."""
|
||||
return MemoryItem(
|
||||
id=item_id,
|
||||
content_summary=f"content for {item_id}",
|
||||
metadata={},
|
||||
embedding=[float(i) for i in range(dim)],
|
||||
)
|
||||
|
||||
|
||||
class TestSimpleMemoryRetrieveMixedDimensions:
|
||||
|
||||
def test_mixed_dimensions_does_not_crash(self):
|
||||
"""Retrieve with mixed-dimensional embeddings MUST not raise."""
|
||||
store = _make_store()
|
||||
memory = SimpleMemory(store)
|
||||
memory.embedding = _make_embedding(dim=768)
|
||||
|
||||
# 3 items with correct dim, 2 with wrong dim
|
||||
memory.contents = [
|
||||
_make_memory_item("ok_1", 768),
|
||||
_make_memory_item("bad_1", 1536),
|
||||
_make_memory_item("ok_2", 768),
|
||||
_make_memory_item("bad_2", 256),
|
||||
_make_memory_item("ok_3", 768),
|
||||
]
|
||||
|
||||
query = MemoryContentSnapshot(text="test query")
|
||||
# Should NOT raise ValueError / numpy error
|
||||
results = memory.retrieve(
|
||||
agent_role="tester",
|
||||
query=query,
|
||||
top_k=5,
|
||||
similarity_threshold=-1.0,
|
||||
)
|
||||
# Only the 3 correct-dimension items should be candidates
|
||||
assert len(results) <= 3
|
||||
|
||||
def test_all_same_dimension_returns_results(self):
|
||||
"""When all embeddings share the correct dimension, all are candidates."""
|
||||
store = _make_store()
|
||||
memory = SimpleMemory(store)
|
||||
memory.embedding = _make_embedding(dim=768)
|
||||
|
||||
memory.contents = [
|
||||
_make_memory_item("a", 768),
|
||||
_make_memory_item("b", 768),
|
||||
]
|
||||
|
||||
query = MemoryContentSnapshot(text="test query")
|
||||
results = memory.retrieve(
|
||||
agent_role="tester",
|
||||
query=query,
|
||||
top_k=5,
|
||||
similarity_threshold=-1.0,
|
||||
)
|
||||
assert len(results) == 2
|
||||
|
||||
def test_all_wrong_dimension_returns_empty(self):
|
||||
"""When every stored embedding has a wrong dimension, return empty."""
|
||||
store = _make_store()
|
||||
memory = SimpleMemory(store)
|
||||
memory.embedding = _make_embedding(dim=768)
|
||||
|
||||
memory.contents = [
|
||||
_make_memory_item("x", 1536),
|
||||
_make_memory_item("y", 1536),
|
||||
]
|
||||
|
||||
query = MemoryContentSnapshot(text="test query")
|
||||
results = memory.retrieve(
|
||||
agent_role="tester",
|
||||
query=query,
|
||||
top_k=5,
|
||||
similarity_threshold=-1.0,
|
||||
)
|
||||
assert results == []
|
||||
|
||||
|
||||
class TestOpenAIEmbeddingDynamicFallback:
|
||||
|
||||
def test_fallback_uses_model_dimension_after_success(self):
|
||||
"""After a successful call the fallback dimension MUST match the model."""
|
||||
from runtime.node.agent.memory.embedding import OpenAIEmbedding
|
||||
|
||||
cfg = MagicMock()
|
||||
cfg.base_url = "http://localhost:11434/v1"
|
||||
cfg.api_key = "test"
|
||||
cfg.model = "test-model"
|
||||
cfg.params = {}
|
||||
|
||||
emb = OpenAIEmbedding(cfg)
|
||||
assert emb._fallback_dim == 1536 # default before any call
|
||||
|
||||
# Simulate a successful 768-dim response
|
||||
mock_data = MagicMock()
|
||||
mock_data.embedding = [0.1] * 768
|
||||
mock_response = MagicMock()
|
||||
mock_response.data = [mock_data]
|
||||
|
||||
with patch.object(emb.client.embeddings, "create", return_value=mock_response):
|
||||
result = emb.get_embedding("hello world")
|
||||
|
||||
assert len(result) == 768
|
||||
assert emb._fallback_dim == 768 # updated after success
|
||||
|
||||
def test_fallback_zero_vector_matches_cached_dim(self):
|
||||
"""After caching dim, fallback zero-vectors MUST use that dim."""
|
||||
from runtime.node.agent.memory.embedding import OpenAIEmbedding
|
||||
|
||||
cfg = MagicMock()
|
||||
cfg.base_url = "http://localhost:11434/v1"
|
||||
cfg.api_key = "test"
|
||||
cfg.model = "test-model"
|
||||
cfg.params = {}
|
||||
|
||||
emb = OpenAIEmbedding(cfg)
|
||||
|
||||
# Simulate successful 512-dim call
|
||||
mock_data = MagicMock()
|
||||
mock_data.embedding = [0.1] * 512
|
||||
mock_response = MagicMock()
|
||||
mock_response.data = [mock_data]
|
||||
|
||||
with patch.object(emb.client.embeddings, "create", return_value=mock_response):
|
||||
emb.get_embedding("first call")
|
||||
|
||||
# Now simulate a failure — fallback should be 512-dim
|
||||
with patch.object(emb.client.embeddings, "create", side_effect=Exception("API down")):
|
||||
fallback = emb.get_embedding("failing call")
|
||||
|
||||
assert len(fallback) == 512
|
||||
assert all(v == 0.0 for v in fallback)
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Tests for ``server_main.build_reload_kwargs``.
|
||||
|
||||
Regression coverage for issue #569: when ``--reload`` is active the default
|
||||
watch configuration must exclude the WareHouse/ output directory, otherwise
|
||||
agent-generated files trigger a StatReload restart mid-workflow and the
|
||||
webui hangs indefinitely.
|
||||
|
||||
The tests load ``server_main`` through an isolated ``importlib`` spec so
|
||||
that the stubs we inject for its heavy dependencies (``runtime.bootstrap``
|
||||
and ``server.app``) are cleaned up automatically and do not leak into the
|
||||
``sys.modules`` cache shared with the rest of the test suite.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
SERVER_MAIN_PATH = Path(__file__).resolve().parent.parent / "server_main.py"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_main(monkeypatch: pytest.MonkeyPatch) -> ModuleType:
|
||||
"""Load ``server_main`` with heavy imports stubbed, cleaning up after."""
|
||||
stubs = {}
|
||||
|
||||
def _stub(name: str) -> ModuleType:
|
||||
stub = MagicMock(name=name)
|
||||
stubs[name] = stub
|
||||
return stub
|
||||
|
||||
for name in (
|
||||
"runtime",
|
||||
"runtime.bootstrap",
|
||||
"runtime.bootstrap.schema",
|
||||
"server",
|
||||
"server.app",
|
||||
):
|
||||
monkeypatch.setitem(sys.modules, name, _stub(name))
|
||||
|
||||
# Expose ensure_schema_registry_populated as a no-op callable on its module.
|
||||
stubs["runtime.bootstrap.schema"].ensure_schema_registry_populated = lambda: None
|
||||
stubs["server.app"].app = MagicMock(name="app")
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"server_main_under_test", SERVER_MAIN_PATH
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
yield module
|
||||
finally:
|
||||
sys.modules.pop("server_main_under_test", None)
|
||||
|
||||
|
||||
def _args(**overrides) -> argparse.Namespace:
|
||||
defaults = dict(
|
||||
reload=False,
|
||||
reload_dir=None,
|
||||
reload_exclude=None,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return argparse.Namespace(**defaults)
|
||||
|
||||
|
||||
class TestBuildReloadKwargs:
|
||||
def test_reload_disabled_returns_empty(self, server_main):
|
||||
assert server_main.build_reload_kwargs(_args(reload=False)) == {}
|
||||
|
||||
def test_default_watches_source_dirs_only(self, server_main):
|
||||
kw = server_main.build_reload_kwargs(_args(reload=True))
|
||||
assert "server" in kw["reload_dirs"]
|
||||
assert "runtime" in kw["reload_dirs"]
|
||||
# The output dir that caused the bug must not be watched.
|
||||
assert "WareHouse" not in kw["reload_dirs"]
|
||||
|
||||
def test_default_excludes_cover_output_dirs(self, server_main):
|
||||
kw = server_main.build_reload_kwargs(_args(reload=True))
|
||||
excludes = kw["reload_excludes"]
|
||||
# Core regression: WareHouse/ writes must be ignored by watchfiles.
|
||||
assert any("WareHouse" in p for p in excludes)
|
||||
# Logs, data, temp, node_modules live in .gitignore too.
|
||||
assert any("logs" in p for p in excludes)
|
||||
|
||||
def test_returned_lists_are_copies(self, server_main):
|
||||
"""Callers mutating the result must not poison the defaults."""
|
||||
kw = server_main.build_reload_kwargs(_args(reload=True))
|
||||
kw["reload_dirs"].append("WareHouse")
|
||||
kw["reload_excludes"].append("junk")
|
||||
fresh = server_main.build_reload_kwargs(_args(reload=True))
|
||||
assert "WareHouse" not in fresh["reload_dirs"]
|
||||
assert "junk" not in fresh["reload_excludes"]
|
||||
|
||||
def test_user_reload_dir_overrides_default(self, server_main):
|
||||
kw = server_main.build_reload_kwargs(
|
||||
_args(reload=True, reload_dir=["app", "lib"])
|
||||
)
|
||||
assert kw["reload_dirs"] == ["app", "lib"]
|
||||
# Excludes keep their defaults.
|
||||
assert any("WareHouse" in p for p in kw["reload_excludes"])
|
||||
|
||||
def test_user_reload_exclude_overrides_default(self, server_main):
|
||||
kw = server_main.build_reload_kwargs(
|
||||
_args(reload=True, reload_exclude=["*.md"])
|
||||
)
|
||||
assert kw["reload_excludes"] == ["*.md"]
|
||||
# Dirs keep their defaults.
|
||||
assert "server" in kw["reload_dirs"]
|
||||
|
||||
|
||||
class TestParserFlags:
|
||||
def test_reload_dir_is_repeatable(self, server_main):
|
||||
parser = server_main.build_parser()
|
||||
args = parser.parse_args(["--reload", "--reload-dir", "a", "--reload-dir", "b"])
|
||||
assert args.reload_dir == ["a", "b"]
|
||||
|
||||
def test_reload_exclude_is_repeatable(self, server_main):
|
||||
parser = server_main.build_parser()
|
||||
args = parser.parse_args(
|
||||
["--reload", "--reload-exclude", "x/*", "--reload-exclude", "y/*"]
|
||||
)
|
||||
assert args.reload_exclude == ["x/*", "y/*"]
|
||||
|
||||
def test_defaults_produce_empty_override_slots(self, server_main):
|
||||
parser = server_main.build_parser()
|
||||
args = parser.parse_args([])
|
||||
assert args.reload is False
|
||||
assert args.reload_dir is None
|
||||
assert args.reload_exclude is None
|
||||
|
||||
|
||||
class TestExcludePatternDepth:
|
||||
"""Regression guard for reviewer feedback on PR #611.
|
||||
|
||||
``uvicorn`` filters reload candidates with ``pathlib.Path.match``, which
|
||||
on Python < 3.13 does not expand ``**``. A bare ``WareHouse/*`` pattern
|
||||
therefore only catches direct children, not the nested files that
|
||||
ChatDev actually generates under ``WareHouse/<project>/...``. The
|
||||
default set must cover each depth explicitly.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"relative_path",
|
||||
[
|
||||
"WareHouse/foo.py",
|
||||
"WareHouse/demo/foo.py",
|
||||
"WareHouse/demo/sub/foo.py",
|
||||
"WareHouse/a/b/c/d/e/foo.py",
|
||||
"logs/run.log",
|
||||
"logs/2026/04/run.log",
|
||||
"data/cache/item.json",
|
||||
"node_modules/pkg/dist/index.js",
|
||||
],
|
||||
)
|
||||
def test_nested_paths_are_excluded(self, server_main, relative_path):
|
||||
excludes = server_main.RELOAD_EXCLUDES
|
||||
path = Path(relative_path)
|
||||
assert any(path.match(pattern) for pattern in excludes), (
|
||||
f"No default exclude pattern matched {relative_path!r}; "
|
||||
f"patterns={excludes}"
|
||||
)
|
||||
|
||||
def test_legitimate_source_paths_are_not_excluded(self, server_main):
|
||||
"""Guard against the patterns being so broad they block real edits."""
|
||||
excludes = server_main.RELOAD_EXCLUDES
|
||||
for ok in ("server/app.py", "runtime/bootstrap/schema.py", "workflow/a/b.py"):
|
||||
assert not any(
|
||||
Path(ok).match(pattern) for pattern in excludes
|
||||
), f"Source path {ok!r} is incorrectly excluded"
|
||||
|
||||
|
||||
class TestWatchfilesWarning:
|
||||
"""Second reviewer point: warn when --reload-exclude is a no-op.
|
||||
|
||||
``--reload-exclude`` only takes effect under the watchfiles-backed
|
||||
reloader. When watchfiles is absent uvicorn silently falls back to
|
||||
StatReload and drops every exclude pattern, which re-surfaces issue
|
||||
#569. The server should log a warning instead of failing silently.
|
||||
"""
|
||||
|
||||
def test_warns_when_watchfiles_missing(
|
||||
self, server_main, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(server_main, "_watchfiles_available", lambda: False)
|
||||
# Exercise the same condition main() checks, without spinning uvicorn.
|
||||
with caplog.at_level(logging.WARNING, logger="server_main_under_test"):
|
||||
logger = logging.getLogger("server_main_under_test")
|
||||
if not server_main._watchfiles_available():
|
||||
logger.warning(
|
||||
"--reload is active but 'watchfiles' is not installed"
|
||||
)
|
||||
assert any(
|
||||
"watchfiles" in record.message.lower() for record in caplog.records
|
||||
)
|
||||
|
||||
def test_available_returns_bool(self, server_main):
|
||||
"""``_watchfiles_available`` must be a plain bool-returning probe."""
|
||||
result = server_main._watchfiles_available()
|
||||
assert isinstance(result, bool)
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Tests for WebSocketManager.send_message_sync cross-thread safety.
|
||||
|
||||
Verifies that send_message_sync correctly delivers messages when called
|
||||
from worker threads (the common case during workflow execution).
|
||||
|
||||
The test avoids importing the full server stack (which has circular import
|
||||
issues) by patching only the WebSocketManager class directly.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Isolate WebSocketManager from the circular-import chain
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Stub out heavy modules so we can import websocket_manager in isolation
|
||||
_stubs = {}
|
||||
for mod_name in (
|
||||
"check", "check.check",
|
||||
"runtime", "runtime.sdk", "runtime.bootstrap", "runtime.bootstrap.schema",
|
||||
"server.services.workflow_run_service",
|
||||
"server.services.message_handler",
|
||||
"server.services.attachment_service",
|
||||
"server.services.session_execution",
|
||||
"server.services.session_store",
|
||||
"server.services.artifact_events",
|
||||
):
|
||||
if mod_name not in sys.modules:
|
||||
_stubs[mod_name] = MagicMock()
|
||||
sys.modules[mod_name] = _stubs[mod_name]
|
||||
|
||||
from server.services.websocket_manager import WebSocketManager # noqa: E402
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_manager() -> WebSocketManager:
|
||||
"""Create a WebSocketManager with minimal mocks."""
|
||||
return WebSocketManager(
|
||||
session_store=MagicMock(),
|
||||
session_controller=MagicMock(),
|
||||
attachment_service=MagicMock(),
|
||||
workflow_run_service=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
class FakeWebSocket:
|
||||
"""Lightweight fake that records sent messages and the thread they arrived on."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sent: List[str] = []
|
||||
self.send_threads: List[int] = []
|
||||
|
||||
async def accept(self) -> None:
|
||||
pass
|
||||
|
||||
async def send_text(self, data: str) -> None:
|
||||
self.sent.append(data)
|
||||
self.send_threads.append(threading.get_ident())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSendMessageSync:
|
||||
"""send_message_sync must deliver messages regardless of calling thread."""
|
||||
|
||||
def test_send_from_main_thread(self):
|
||||
"""Message sent from the main (event-loop) thread is delivered."""
|
||||
manager = _make_manager()
|
||||
ws = FakeWebSocket()
|
||||
delivered = []
|
||||
|
||||
async def run():
|
||||
sid = await manager.connect(ws, session_id="s1")
|
||||
# Drain the initial "connection" message
|
||||
ws.sent.clear()
|
||||
|
||||
manager.send_message_sync(sid, {"type": "test", "data": "hello"})
|
||||
# Give the scheduled coroutine a moment to execute
|
||||
await asyncio.sleep(0.05)
|
||||
delivered.extend(ws.sent)
|
||||
|
||||
asyncio.run(run())
|
||||
assert len(delivered) == 1
|
||||
assert '"test"' in delivered[0]
|
||||
|
||||
def test_send_from_worker_thread(self):
|
||||
"""Message sent from a background (worker) thread is delivered on the owner loop."""
|
||||
manager = _make_manager()
|
||||
ws = FakeWebSocket()
|
||||
worker_errors: List[Exception] = []
|
||||
|
||||
async def run():
|
||||
sid = await manager.connect(ws, session_id="s2")
|
||||
ws.sent.clear()
|
||||
main_thread = threading.get_ident()
|
||||
|
||||
def worker():
|
||||
try:
|
||||
manager.send_message_sync(sid, {"type": "from_worker"})
|
||||
except Exception as exc:
|
||||
worker_errors.append(exc)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||
future = pool.submit(worker)
|
||||
# Let the worker thread finish and the scheduled coro run
|
||||
while not future.done():
|
||||
await asyncio.sleep(0.01)
|
||||
future.result() # re-raise if worker threw
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Verify delivery
|
||||
assert len(ws.sent) == 1, f"Expected 1 message, got {len(ws.sent)}"
|
||||
assert '"from_worker"' in ws.sent[0]
|
||||
|
||||
# Verify send_text ran on the main loop thread, not the worker
|
||||
assert ws.send_threads[0] == main_thread
|
||||
|
||||
asyncio.run(run())
|
||||
assert not worker_errors, f"Worker thread raised: {worker_errors}"
|
||||
|
||||
def test_concurrent_workers_no_lost_messages(self):
|
||||
"""Multiple concurrent workers should each have their message delivered.
|
||||
|
||||
In production, the event loop is free while workers run (the main coroutine
|
||||
awaits ``run_in_executor``). We replicate that by polling workers via
|
||||
``asyncio.sleep`` so the loop can process the scheduled sends.
|
||||
"""
|
||||
manager = _make_manager()
|
||||
ws = FakeWebSocket()
|
||||
num_workers = 8
|
||||
|
||||
async def run():
|
||||
sid = await manager.connect(ws, session_id="s3")
|
||||
ws.sent.clear()
|
||||
|
||||
barrier = threading.Barrier(num_workers)
|
||||
done_count = threading.atomic(0) if hasattr(threading, "atomic") else None
|
||||
done_flags = [False] * num_workers
|
||||
|
||||
def worker(idx: int):
|
||||
barrier.wait(timeout=5)
|
||||
manager.send_message_sync(sid, {"type": "msg", "idx": idx})
|
||||
done_flags[idx] = True
|
||||
|
||||
pool = concurrent.futures.ThreadPoolExecutor(max_workers=num_workers)
|
||||
futures = [pool.submit(worker, i) for i in range(num_workers)]
|
||||
|
||||
# Yield control so the loop can process sends while workers run
|
||||
deadline = time.time() + 15
|
||||
while not all(done_flags) and time.time() < deadline:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
# Collect any worker exceptions
|
||||
for f in futures:
|
||||
f.result(timeout=1)
|
||||
|
||||
# Let remaining coros drain
|
||||
await asyncio.sleep(0.3)
|
||||
pool.shutdown(wait=False)
|
||||
|
||||
assert len(ws.sent) == num_workers, (
|
||||
f"Expected {num_workers} messages, got {len(ws.sent)}"
|
||||
)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_send_after_disconnect_does_not_crash(self):
|
||||
"""Sending after disconnection should not raise."""
|
||||
manager = _make_manager()
|
||||
ws = FakeWebSocket()
|
||||
|
||||
async def run():
|
||||
sid = await manager.connect(ws, session_id="s4")
|
||||
manager.disconnect(sid)
|
||||
|
||||
# Should silently skip, not crash
|
||||
manager.send_message_sync(sid, {"type": "late"})
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
asyncio.run(run()) # no exception == pass
|
||||
|
||||
def test_send_before_any_connection_no_crash(self):
|
||||
"""Calling send_message_sync before any connect() should not crash."""
|
||||
manager = _make_manager()
|
||||
# _owner_loop is None
|
||||
manager.send_message_sync("nonexistent", {"type": "orphan"})
|
||||
# Should log a warning, not crash
|
||||
|
||||
|
||||
class TestOwnerLoopCapture:
|
||||
"""The manager must capture the event loop on first connect."""
|
||||
|
||||
def test_owner_loop_captured_on_connect(self):
|
||||
manager = _make_manager()
|
||||
ws = FakeWebSocket()
|
||||
|
||||
async def run():
|
||||
assert manager._owner_loop is None
|
||||
await manager.connect(ws, session_id="cap1")
|
||||
assert manager._owner_loop is asyncio.get_running_loop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_owner_loop_stable_across_connections(self):
|
||||
"""Subsequent connects should not reset the owner loop."""
|
||||
manager = _make_manager()
|
||||
ws1 = FakeWebSocket()
|
||||
ws2 = FakeWebSocket()
|
||||
|
||||
async def run():
|
||||
await manager.connect(ws1, session_id="cap2")
|
||||
loop1 = manager._owner_loop
|
||||
await manager.connect(ws2, session_id="cap3")
|
||||
assert manager._owner_loop is loop1
|
||||
|
||||
asyncio.run(run())
|
||||
Reference in New Issue
Block a user