mirror of
https://github.com/langgenius/dify.git
synced 2026-05-26 22:00:53 -04:00
55f95dbc36
Co-authored-by: Copilot <copilot@github.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
402 lines
13 KiB
Python
402 lines
13 KiB
Python
import asyncio
|
|
from collections.abc import Iterator
|
|
from dataclasses import dataclass, field
|
|
from itertools import count
|
|
|
|
import pytest
|
|
from pydantic import BaseModel, ConfigDict, ValidationError
|
|
from typing_extensions import override
|
|
|
|
from agenton.compositor import Compositor, CompositorSessionSnapshot, LayerNode, LayerProvider
|
|
from agenton.layers import (
|
|
EmptyLayerConfig,
|
|
ExitIntent,
|
|
LayerConfig,
|
|
LifecycleState,
|
|
NoLayerDeps,
|
|
PlainLayer,
|
|
)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TraceLayer(PlainLayer[NoLayerDeps]):
|
|
"""Layer that records no-arg lifecycle events observable to tests."""
|
|
|
|
events: list[str] = field(default_factory=list)
|
|
|
|
@override
|
|
async def on_context_create(self) -> None:
|
|
self.events.append("create")
|
|
|
|
@override
|
|
async def on_context_suspend(self) -> None:
|
|
self.events.append("suspend")
|
|
|
|
@override
|
|
async def on_context_resume(self) -> None:
|
|
self.events.append("resume")
|
|
|
|
@override
|
|
async def on_context_delete(self) -> None:
|
|
self.events.append("delete")
|
|
|
|
|
|
def _compositor(*layer_names: str) -> Compositor:
|
|
return Compositor([LayerNode(layer_name, TraceLayer) for layer_name in layer_names])
|
|
|
|
|
|
def test_same_compositor_enters_multiple_times_with_fresh_layers_and_snapshot_resume() -> None:
|
|
compositor = _compositor("first", "second")
|
|
runs = []
|
|
|
|
async def run() -> None:
|
|
async with compositor.enter() as first_run:
|
|
assert [slot.lifecycle_state for slot in first_run.slots.values()] == [
|
|
LifecycleState.ACTIVE,
|
|
LifecycleState.ACTIVE,
|
|
]
|
|
first_run.suspend_on_exit()
|
|
assert [slot.exit_intent for slot in first_run.slots.values()] == [
|
|
ExitIntent.SUSPEND,
|
|
ExitIntent.SUSPEND,
|
|
]
|
|
runs.append(first_run)
|
|
|
|
assert first_run.session_snapshot is not None
|
|
async with compositor.enter(session_snapshot=first_run.session_snapshot) as resumed_run:
|
|
assert resumed_run.get_layer("first", TraceLayer).events == ["resume"]
|
|
assert resumed_run.get_layer("second", TraceLayer).events == ["resume"]
|
|
runs.append(resumed_run)
|
|
|
|
asyncio.run(run())
|
|
|
|
first_layer = runs[0].get_layer("first", TraceLayer)
|
|
resumed_layer = runs[1].get_layer("first", TraceLayer)
|
|
assert first_layer is not resumed_layer
|
|
assert first_layer.events == ["create", "suspend"]
|
|
assert resumed_layer.events == ["resume", "delete"]
|
|
assert runs[1].session_snapshot is not None
|
|
assert [layer.lifecycle_state for layer in runs[1].session_snapshot.layers] == [
|
|
LifecycleState.CLOSED,
|
|
LifecycleState.CLOSED,
|
|
]
|
|
|
|
|
|
def test_concurrent_enters_do_not_share_layer_instances() -> None:
|
|
compositor = _compositor("trace")
|
|
|
|
async def enter_once() -> tuple[int, list[str]]:
|
|
async with compositor.enter() as run:
|
|
layer = run.get_layer("trace", TraceLayer)
|
|
await asyncio.sleep(0)
|
|
return id(layer), layer.events
|
|
|
|
async def run_concurrently() -> list[tuple[int, list[str]]]:
|
|
return list(await asyncio.gather(enter_once(), enter_once()))
|
|
|
|
results = asyncio.run(run_concurrently())
|
|
|
|
assert results[0][0] != results[1][0]
|
|
assert results[0][1] == ["create", "delete"]
|
|
assert results[1][1] == ["create", "delete"]
|
|
|
|
|
|
class ConfiguredLayerConfig(LayerConfig):
|
|
value: str
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ConfiguredLayer(PlainLayer[NoLayerDeps, ConfiguredLayerConfig]):
|
|
type_id = "test.configured"
|
|
|
|
value: str
|
|
|
|
hooks: list[str] = field(default_factory=list)
|
|
|
|
@classmethod
|
|
@override
|
|
def from_config(cls, config: ConfiguredLayerConfig) -> "ConfiguredLayer":
|
|
return cls(value=config.value)
|
|
|
|
@override
|
|
async def on_context_create(self) -> None:
|
|
self.hooks.append(f"create:{self.config.value}")
|
|
|
|
|
|
def test_custom_factory_is_called_each_enter_with_typed_config() -> None:
|
|
calls: list[str] = []
|
|
|
|
def create_layer(config: ConfiguredLayerConfig) -> ConfiguredLayer:
|
|
calls.append(config.value)
|
|
return ConfiguredLayer(value=f"factory:{config.value}")
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("configured", LayerProvider.from_factory(layer_type=ConfiguredLayer, create=create_layer))]
|
|
)
|
|
|
|
async def run() -> None:
|
|
async with compositor.enter(configs={"configured": {"value": "one"}}) as first_run:
|
|
first_layer = first_run.get_layer("configured", ConfiguredLayer)
|
|
assert first_layer.value == "factory:one"
|
|
assert first_layer.config.value == "one"
|
|
async with compositor.enter(configs={"configured": ConfiguredLayerConfig(value="two")}) as second_run:
|
|
second_layer = second_run.get_layer("configured", ConfiguredLayer)
|
|
assert second_layer.value == "factory:two"
|
|
assert second_layer.config.value == "two"
|
|
assert second_layer is not first_layer
|
|
|
|
asyncio.run(run())
|
|
|
|
assert calls == ["one", "two"]
|
|
|
|
|
|
def test_provider_rejects_reused_layer_instance_before_hooks_run() -> None:
|
|
shared_layer = TraceLayer()
|
|
compositor = Compositor(
|
|
[
|
|
LayerNode(
|
|
"trace",
|
|
LayerProvider.from_factory(layer_type=TraceLayer, create=lambda config: shared_layer),
|
|
)
|
|
]
|
|
)
|
|
|
|
async def run() -> None:
|
|
async with compositor.enter():
|
|
pass
|
|
|
|
with pytest.raises(ValueError, match="fresh layer instance"):
|
|
async with compositor.enter():
|
|
pass
|
|
|
|
asyncio.run(run())
|
|
|
|
assert shared_layer.events == ["create", "delete"]
|
|
|
|
|
|
def test_configs_are_validated_by_node_name_before_factory_call() -> None:
|
|
calls: list[str] = []
|
|
|
|
def create_layer(config: ConfiguredLayerConfig) -> ConfiguredLayer:
|
|
calls.append(config.value)
|
|
return ConfiguredLayer(value=config.value)
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("configured", LayerProvider.from_factory(layer_type=ConfiguredLayer, create=create_layer))]
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="unknown layer node names: missing"):
|
|
asyncio.run(_enter_once(compositor, configs={"missing": {}}))
|
|
|
|
with pytest.raises(ValidationError):
|
|
asyncio.run(_enter_once(compositor, configs={"configured": {"unknown": "field"}}))
|
|
|
|
assert calls == []
|
|
|
|
|
|
def test_all_node_configs_are_validated_before_any_factory_runs() -> None:
|
|
calls: list[str] = []
|
|
|
|
def create_layer(config: ConfiguredLayerConfig) -> ConfiguredLayer:
|
|
calls.append(config.value)
|
|
return ConfiguredLayer(value=config.value)
|
|
|
|
provider = LayerProvider.from_factory(layer_type=ConfiguredLayer, create=create_layer)
|
|
compositor = Compositor([LayerNode("first", provider), LayerNode("second", provider)])
|
|
|
|
with pytest.raises(ValidationError):
|
|
asyncio.run(
|
|
_enter_once(
|
|
compositor,
|
|
configs={"first": {"value": "valid"}, "second": {"unknown": "field"}},
|
|
)
|
|
)
|
|
|
|
assert calls == []
|
|
|
|
|
|
def test_existing_config_model_instances_are_revalidated_before_factory_runs() -> None:
|
|
calls: list[str] = []
|
|
|
|
def create_layer(config: ConfiguredLayerConfig) -> ConfiguredLayer:
|
|
calls.append(config.value)
|
|
return ConfiguredLayer(value=config.value)
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("configured", LayerProvider.from_factory(layer_type=ConfiguredLayer, create=create_layer))]
|
|
)
|
|
config = ConfiguredLayerConfig(value="valid")
|
|
config.value = 123 # pyright: ignore[reportAttributeAccessIssue]
|
|
|
|
with pytest.raises(ValidationError):
|
|
asyncio.run(_enter_once(compositor, configs={"configured": config}))
|
|
|
|
assert calls == []
|
|
|
|
|
|
def test_existing_snapshot_model_instances_are_revalidated_before_factory_runs() -> None:
|
|
calls = 0
|
|
|
|
def create_layer(config: EmptyLayerConfig) -> TraceLayer:
|
|
nonlocal calls
|
|
calls += 1
|
|
return TraceLayer()
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("trace", LayerProvider.from_factory(layer_type=TraceLayer, create=create_layer))]
|
|
)
|
|
snapshot = CompositorSessionSnapshot.model_validate(
|
|
{"layers": [{"name": "trace", "lifecycle_state": "suspended", "runtime_state": {}}]}
|
|
)
|
|
snapshot.layers[0].lifecycle_state = LifecycleState.ACTIVE
|
|
|
|
with pytest.raises(ValidationError, match="ACTIVE is internal-only"):
|
|
asyncio.run(_enter_once(compositor, session_snapshot=snapshot))
|
|
|
|
assert calls == 0
|
|
|
|
|
|
class RuntimeState(BaseModel):
|
|
runtime_id: int | None = None
|
|
resumed_runtime_id: int | None = None
|
|
deleted_runtime_id: int | None = None
|
|
body_value: str | None = None
|
|
|
|
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RuntimeStateLayer(PlainLayer[NoLayerDeps, EmptyLayerConfig, RuntimeState]):
|
|
next_id: Iterator[int] = field(default_factory=lambda: count(1))
|
|
|
|
@override
|
|
async def on_context_create(self) -> None:
|
|
self.runtime_state.runtime_id = next(self.next_id)
|
|
|
|
@override
|
|
async def on_context_resume(self) -> None:
|
|
self.runtime_state.resumed_runtime_id = self.runtime_state.runtime_id
|
|
|
|
@override
|
|
async def on_context_delete(self) -> None:
|
|
self.runtime_state.deleted_runtime_id = self.runtime_state.runtime_id
|
|
|
|
|
|
def test_snapshot_hydrates_runtime_state_and_exit_snapshots_from_layer_self() -> None:
|
|
compositor = Compositor([LayerNode("state", RuntimeStateLayer)])
|
|
|
|
async def create_suspend_resume_delete() -> tuple[CompositorSessionSnapshot, CompositorSessionSnapshot]:
|
|
async with compositor.enter() as first_run:
|
|
first_run.suspend_on_exit()
|
|
assert first_run.session_snapshot is not None
|
|
|
|
async with compositor.enter(session_snapshot=first_run.session_snapshot) as resumed_run:
|
|
resumed_layer = resumed_run.get_layer("state", RuntimeStateLayer)
|
|
assert isinstance(resumed_layer.runtime_state, RuntimeState)
|
|
assert resumed_layer.runtime_state.runtime_id == 1
|
|
assert resumed_layer.runtime_state.resumed_runtime_id == 1
|
|
resumed_layer.runtime_state.body_value = "mutated on self"
|
|
assert resumed_run.session_snapshot is not None
|
|
return first_run.session_snapshot, resumed_run.session_snapshot
|
|
|
|
suspended_snapshot, closed_snapshot = asyncio.run(create_suspend_resume_delete())
|
|
|
|
assert suspended_snapshot.model_dump(mode="json") == {
|
|
"schema_version": 1,
|
|
"layers": [
|
|
{
|
|
"name": "state",
|
|
"lifecycle_state": "suspended",
|
|
"runtime_state": {
|
|
"runtime_id": 1,
|
|
"resumed_runtime_id": None,
|
|
"deleted_runtime_id": None,
|
|
"body_value": None,
|
|
},
|
|
}
|
|
],
|
|
}
|
|
assert closed_snapshot.model_dump(mode="json") == {
|
|
"schema_version": 1,
|
|
"layers": [
|
|
{
|
|
"name": "state",
|
|
"lifecycle_state": "closed",
|
|
"runtime_state": {
|
|
"runtime_id": 1,
|
|
"resumed_runtime_id": 1,
|
|
"deleted_runtime_id": 1,
|
|
"body_value": "mutated on self",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def test_run_snapshot_rejects_active_layers() -> None:
|
|
compositor = _compositor("trace")
|
|
|
|
async def run() -> None:
|
|
async with compositor.enter() as active_run:
|
|
with pytest.raises(RuntimeError, match="Cannot snapshot active compositor run layers: trace"):
|
|
active_run.snapshot_session()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_active_snapshot_input_is_rejected_before_factories_run() -> None:
|
|
calls = 0
|
|
|
|
def create_layer(config: EmptyLayerConfig) -> TraceLayer:
|
|
nonlocal calls
|
|
calls += 1
|
|
return TraceLayer()
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("trace", LayerProvider.from_factory(layer_type=TraceLayer, create=create_layer))]
|
|
)
|
|
active_snapshot = {"layers": [{"name": "trace", "lifecycle_state": "active", "runtime_state": {}}]}
|
|
|
|
with pytest.raises(ValidationError, match="ACTIVE is internal-only"):
|
|
CompositorSessionSnapshot.model_validate(active_snapshot)
|
|
|
|
with pytest.raises(ValidationError, match="ACTIVE is internal-only"):
|
|
asyncio.run(_enter_once(compositor, session_snapshot=active_snapshot))
|
|
|
|
assert calls == 0
|
|
|
|
|
|
def test_closed_snapshot_enter_is_rejected_before_hooks_run() -> None:
|
|
created_layers: list[TraceLayer] = []
|
|
|
|
def create_layer(config: EmptyLayerConfig) -> TraceLayer:
|
|
layer = TraceLayer()
|
|
created_layers.append(layer)
|
|
return layer
|
|
|
|
compositor = Compositor(
|
|
[LayerNode("trace", LayerProvider.from_factory(layer_type=TraceLayer, create=create_layer))]
|
|
)
|
|
closed_snapshot = {"layers": [{"name": "trace", "lifecycle_state": "closed", "runtime_state": {}}]}
|
|
|
|
with pytest.raises(RuntimeError, match="CLOSED snapshots cannot be entered"):
|
|
asyncio.run(_enter_once(compositor, session_snapshot=closed_snapshot))
|
|
|
|
assert len(created_layers) == 1
|
|
assert created_layers[0].events == []
|
|
|
|
|
|
async def _enter_once(
|
|
compositor: Compositor,
|
|
*,
|
|
configs: dict[str, object] | None = None,
|
|
session_snapshot: object | None = None,
|
|
) -> None:
|
|
async with compositor.enter(
|
|
configs=configs, # pyright: ignore[reportArgumentType]
|
|
session_snapshot=session_snapshot, # pyright: ignore[reportArgumentType]
|
|
):
|
|
pass
|