feat(agentscope): add memory system and automation job support

- Add consumer_registry and pipeline_registry for runtime orchestration
- Add Visibility schema for message filtering
- Add PipelineSpec for agent pipeline configuration
- Add automation job models and configuration
- Remove memory_prompt.py (consolidated into memory system)
- Update runtime components: context_loader, context_service, orchestrator, runner, tasks
- Update toolkit: tool_config, tool_middleware, custom tools (calendar, user_lookup)
- Add auth_helpers and calendar_domain utilities
- Add system_agents.yaml configuration
This commit is contained in:
qzl
2026-03-19 18:42:35 +08:00
parent 0661016827
commit 0abf51e837
55 changed files with 2172 additions and 1233 deletions
@@ -189,3 +189,56 @@ def test_step_started_internal_event_keeps_step_name() -> None:
assert result["type"] == "STEP_STARTED"
assert result["stepName"] == "worker"
def test_run_error_prefers_top_level_message_and_code() -> None:
internal = {
"type": "run.error",
"threadId": "thread-1",
"runId": "run-1",
"message": "runtime failed",
"code": "RUNTIME_ERROR",
"data": {
"message": "nested message",
"code": "NESTED_ERROR",
},
}
result = to_agui_wire_event(internal)
assert result["type"] == "RUN_ERROR"
assert result["message"] == "runtime failed"
assert result["code"] == "RUNTIME_ERROR"
def test_run_error_falls_back_to_data_when_top_level_missing() -> None:
internal = {
"type": "run.error",
"threadId": "thread-1",
"runId": "run-1",
"data": {
"message": "nested message",
"code": "NESTED_ERROR",
},
}
result = to_agui_wire_event(internal)
assert result["type"] == "RUN_ERROR"
assert result["message"] == "nested message"
assert result["code"] == "NESTED_ERROR"
def test_run_error_uses_default_message_when_payload_invalid() -> None:
internal = {
"type": "run.error",
"threadId": "thread-1",
"runId": "run-1",
"data": "invalid",
}
result = to_agui_wire_event(internal)
assert result["type"] == "RUN_ERROR"
assert result["message"] == "Unknown error"
assert "code" not in result
@@ -59,6 +59,16 @@ def _patch_repositories(
monkeypatch.setattr(store_module, "MessageRepository", _FakeMessageRepository)
monkeypatch.setattr(store_module, "AgentChatSessionStatus", _SessionStatus)
async def _fake_stage_bit_map(self, *, session: object) -> dict[str, int]:
del self, session
return {"router": 16, "worker": 17, "memory": 18}
monkeypatch.setattr(
store_module.SqlAlchemyEventStore,
"_load_stage_visibility_bit_map",
_fake_stage_bit_map,
)
@pytest.mark.asyncio
async def test_store_persists_worker_output_with_answer_as_content(
@@ -103,6 +113,7 @@ async def test_store_persists_worker_output_with_answer_as_content(
assert metadata["agent_output"]["answer"] == "worker-answer"
assert metadata["agent_output"]["ui_hints"]["intent"] == "message"
assert append_kwargs["cost"] == Decimal("0.123")
assert append_kwargs["visibility_mask"] == ((1 << 0) | (1 << 17))
assert captured["message_delta"] == 1
assert captured["token_delta"] == 8
@@ -141,3 +152,4 @@ async def test_store_persists_tool_output_with_summary_as_content(
metadata["tool_agent_output"]["result"]
== "status=success batch=1 success=1 failed=0 ids=[event-1]"
)
assert append_kwargs["visibility_mask"] == (1 << 0)
@@ -1,12 +1,13 @@
from __future__ import annotations
from uuid import UUID
from uuid import uuid4
import pytest
from core.agentscope.persistence.user_context_cache import UserContextCache
from core.agentscope.schemas.user_context import (
UserAgentContext,
from schemas.user.context import (
UserContext,
parse_profile_settings,
)
@@ -45,15 +46,15 @@ class _FakeRedis:
self.set_store.pop(key, None)
return len(keys)
async def sadd(self, key: str, *values: str) -> int:
bucket = self.set_store.setdefault(key, set())
async def sadd(self, name: str, *values: str) -> int:
bucket = self.set_store.setdefault(name, set())
before = len(bucket)
for value in values:
bucket.add(value)
return len(bucket) - before
async def smembers(self, key: str) -> set[str]:
return set(self.set_store.get(key, set()))
async def smembers(self, name: str) -> set[str]:
return set(self.set_store.get(name, set()))
class _BrokenRedis:
@@ -77,18 +78,18 @@ class _BrokenRedis:
del keys
raise RuntimeError("redis down")
async def sadd(self, key: str, *values: str) -> int:
del key, values
async def sadd(self, name: str, *values: str) -> int:
del name, values
raise RuntimeError("redis down")
async def smembers(self, key: str) -> set[str]:
del key
async def smembers(self, name: str) -> set[str]:
del name
raise RuntimeError("redis down")
def _build_context() -> UserAgentContext:
return UserAgentContext(
user_id=uuid4(),
def _build_context() -> UserContext:
return UserContext(
id=str(uuid4()),
username="demo-user",
bio="demo bio",
settings=parse_profile_settings({"preferences": {"ai_language": "en-US"}}),
@@ -111,11 +112,11 @@ async def test_user_context_cache_set_and_get_hit() -> None:
loaded = await cache.get(session_id=session_id)
assert loaded is not None
assert loaded.user_id == context.user_id
assert loaded.id == context.id
assert loaded.username == "demo-user"
assert redis.expire_calls == [
(f"agent:user-context:{session_id}", 600),
(f"agent:user-context:sessions:{context.user_id}", 600),
(f"agent:user-context:sessions:{context.id}", 600),
]
assert redis.hincrby_calls == [
(f"agent:user-context:{session_id}", "turns_used", 1)
@@ -138,12 +139,14 @@ async def test_user_context_cache_invalidate_user_deletes_all_sessions() -> None
await cache.set(session_id=s1, context=context)
await cache.set(session_id=s2, context=context)
deleted = await cache.invalidate_user(user_id=context.user_id)
deleted = await cache.invalidate_user(user_id=UUID(context.id))
assert deleted == 2
assert f"agent:user-context:{s1}" in redis.delete_calls
assert f"agent:user-context:{s2}" in redis.delete_calls
assert f"agent:user-context:sessions:{context.user_id}" in redis.delete_calls
assert any(
key.startswith("agent:user-context:sessions:") for key in redis.delete_calls
)
@pytest.mark.asyncio
@@ -0,0 +1,28 @@
from __future__ import annotations
import pytest
from core.agentscope.runtime.consumer_registry import build_consumer_registry
def test_build_consumer_registry_from_system_agent_configs() -> None:
registry = build_consumer_registry(
system_agent_configs={
"router": {"config": {"visibility_consumer_bit": 16}},
"worker": {"config": {"visibility_consumer_bit": 17}},
"memory": {"config": {"visibility_consumer_bit": 18}},
}
)
assert registry.resolve_agent_bit(agent_type="router") == 16
assert registry.resolve_agent_bit(agent_type="worker") == 17
def test_build_consumer_registry_rejects_duplicate_bit() -> None:
with pytest.raises(ValueError, match="duplicate visibility bit"):
build_consumer_registry(
system_agent_configs={
"router": {"config": {"visibility_consumer_bit": 16}},
"worker": {"config": {"visibility_consumer_bit": 16}},
}
)
@@ -28,7 +28,7 @@ def _user_context() -> UserContext:
return UserContext(
id="00000000-0000-0000-0000-000000000001",
username="alice",
email="alice@example.com",
phone="+8613900000000",
settings=parse_profile_settings(None),
)
@@ -42,7 +42,7 @@ def _run_input() -> RunAgentInput:
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
"tools": [],
"context": [],
"forwardedProps": {},
"forwardedProps": {"agent_type": "worker"},
}
)
@@ -58,6 +58,7 @@ async def test_orchestrator_emits_run_lifecycle_events() -> None:
run_input=_run_input(),
context_messages=[],
user_context=_user_context(),
system_agent_mode="worker",
)
assert result["worker"]["answer"] == "done"
@@ -0,0 +1,24 @@
from __future__ import annotations
import pytest
from core.agentscope.runtime.pipeline_registry import build_default_pipeline_spec
def test_build_default_pipeline_spec_worker_has_two_stages() -> None:
spec = build_default_pipeline_spec(mode="worker")
assert spec.mode == "worker"
assert [item.stage_name for item in spec.stages] == ["router", "worker"]
def test_build_default_pipeline_spec_memory_has_single_stage() -> None:
spec = build_default_pipeline_spec(mode="memory")
assert spec.mode == "memory"
assert [item.stage_name for item in spec.stages] == ["memory"]
def test_build_default_pipeline_spec_rejects_unknown_mode() -> None:
with pytest.raises(ValueError, match="unsupported pipeline mode"):
build_default_pipeline_spec(mode="planner")
@@ -3,8 +3,23 @@ from __future__ import annotations
import pytest
from ag_ui.core import RunAgentInput
import core.agentscope.runtime.runner as runner_module
from core.agentscope.runtime.runner import AgentScopeRunner
from schemas.automation.config import default_memory_job_config
from schemas.agent.runtime_models import (
ExecutionMode,
NormalizedTaskInput,
ResultType,
ResultTyping,
RouterAgentOutput,
RouterUiDecision,
TaskType,
TaskTyping,
UiMode,
WorkerAgentOutputLite,
)
from schemas.agent.system_agent import AgentType
from schemas.user import UserContext, parse_profile_settings
def _run_input() -> RunAgentInput:
@@ -16,19 +31,51 @@ def _run_input() -> RunAgentInput:
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
"tools": [],
"context": [],
"forwardedProps": {},
"forwardedProps": {"agent_type": "worker"},
}
)
def test_resolve_stage_agent_type_defaults_to_worker() -> None:
assert AgentScopeRunner._resolve_stage_agent_type("") == AgentType.WORKER
assert AgentScopeRunner._resolve_stage_agent_type("worker") == AgentType.WORKER
assert AgentScopeRunner._resolve_stage_agent_type("unknown") == AgentType.WORKER
def _user_context() -> UserContext:
return UserContext(
id="00000000-0000-0000-0000-000000000001",
username="alice",
phone="+8613900000000",
settings=parse_profile_settings(None),
)
def test_resolve_stage_agent_type_supports_memory() -> None:
assert AgentScopeRunner._resolve_stage_agent_type("memory") == AgentType.MEMORY
def test_parse_agent_type_supports_known_stages() -> None:
assert AgentScopeRunner._parse_agent_type(stage_name="router") == AgentType.ROUTER
assert AgentScopeRunner._parse_agent_type(stage_name="worker") == AgentType.WORKER
assert AgentScopeRunner._parse_agent_type(stage_name="memory") == AgentType.MEMORY
def test_parse_agent_type_rejects_unknown_stage() -> None:
with pytest.raises(ValueError, match="unsupported stage name"):
AgentScopeRunner._parse_agent_type(stage_name="planner")
def test_build_worker_input_messages_only_contains_router_contract() -> None:
runner = AgentScopeRunner()
router_output = RouterAgentOutput(
normalized_task_input=NormalizedTaskInput(user_text="安排明天会议"),
key_entities=[],
constraints=[],
task_typing=TaskTyping(primary=TaskType.SCHEDULING),
execution_mode=ExecutionMode.TOOL_ASSISTED,
result_typing=ResultTyping(primary=ResultType.EXECUTION_REPORT),
ui=RouterUiDecision(
ui_mode=UiMode.NONE,
ui_decision_reason="单一执行任务,文本输出足够",
),
)
input_messages = runner._build_worker_input_messages(router_output=router_output)
assert len(input_messages) == 1
assert input_messages[0].role == "user"
assert "[RouterAgentOutput]" in str(input_messages[0].content)
@pytest.mark.asyncio
@@ -43,11 +90,12 @@ async def test_resolve_runtime_client_time_from_forwarded_props() -> None:
"tools": [],
"context": [],
"forwardedProps": {
"agent_type": "worker",
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16T09:12:33-07:00",
"client_epoch_ms": 1773658353000,
}
},
},
}
)
@@ -55,3 +103,146 @@ async def test_resolve_runtime_client_time_from_forwarded_props() -> None:
resolved = runner._resolve_runtime_client_time(run_input=run_input)
assert resolved is not None
assert resolved.device_timezone == "America/Los_Angeles"
@pytest.mark.asyncio
async def test_execute_worker_mode_runs_router_then_worker(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakePipeline:
async def emit(self, *, session_id: str, event: dict[str, object]) -> str:
del session_id, event
return "1-0"
class _FakeSessionCtx:
async def __aenter__(self) -> object:
return object()
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
del exc_type, exc, tb
runner = AgentScopeRunner()
load_calls: list[AgentType] = []
async def _fake_load_stage_config(*, session: object, agent_type: AgentType):
del session
load_calls.append(agent_type)
return runner_module.SystemAgentRuntimeConfig(
agent_type=agent_type,
model_code="demo",
api_base_url="https://example.com",
api_key="test",
llm_config=runner_module.SystemAgentLLMConfig(),
)
async def _fake_execute_router_step(**kwargs: object) -> RouterAgentOutput:
del kwargs
return RouterAgentOutput(
normalized_task_input=NormalizedTaskInput(user_text="安排会议"),
key_entities=[],
constraints=[],
task_typing=TaskTyping(primary=TaskType.SCHEDULING),
execution_mode=ExecutionMode.TOOL_ASSISTED,
result_typing=ResultTyping(primary=ResultType.EXECUTION_REPORT),
ui=RouterUiDecision(
ui_mode=UiMode.NONE,
ui_decision_reason="单任务",
),
)
async def _fake_execute_worker_step(**kwargs: object) -> WorkerAgentOutputLite:
del kwargs
return WorkerAgentOutputLite(answer="ok")
monkeypatch.setattr(runner_module, "AsyncSessionLocal", lambda: _FakeSessionCtx())
monkeypatch.setattr(runner, "_load_stage_config", _fake_load_stage_config)
monkeypatch.setattr(runner, "_build_stage_toolkit", lambda **kwargs: object())
monkeypatch.setattr(runner, "_execute_router_step", _fake_execute_router_step)
monkeypatch.setattr(runner, "_execute_worker_step", _fake_execute_worker_step)
result = await runner.execute(
user_context=_user_context(),
context_messages=[],
pipeline=_FakePipeline(),
run_input=_run_input(),
system_agent_mode="worker",
)
assert load_calls == [AgentType.ROUTER, AgentType.WORKER]
assert result["router"]["normalized_task_input"]["user_text"] == "安排会议"
assert result["worker"]["answer"] == "ok"
@pytest.mark.asyncio
async def test_execute_memory_mode_requires_memory_job_config() -> None:
runner = AgentScopeRunner()
class _FakePipeline:
async def emit(self, *, session_id: str, event: dict[str, object]) -> str:
del session_id, event
return "1-0"
with pytest.raises(RuntimeError, match="memory job config is required"):
await runner.execute(
user_context=_user_context(),
context_messages=[],
pipeline=_FakePipeline(),
run_input=_run_input(),
system_agent_mode="memory",
memory_job_config=None,
)
@pytest.mark.asyncio
async def test_execute_memory_mode_uses_memory_job_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakePipeline:
async def emit(self, *, session_id: str, event: dict[str, object]) -> str:
del session_id, event
return "1-0"
class _FakeSessionCtx:
async def __aenter__(self) -> object:
return object()
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
del exc_type, exc, tb
runner = AgentScopeRunner()
async def _fake_build_memory_stage_config(**kwargs: object):
del kwargs
return runner_module.SystemAgentRuntimeConfig(
agent_type=AgentType.MEMORY,
model_code="qwen3.5-flash",
api_base_url="https://example.com",
api_key="test",
llm_config=runner_module.SystemAgentLLMConfig(),
)
async def _fake_execute_single_stage_step(**kwargs: object):
del kwargs
return runner_module.AgentOutput(answer="memory")
monkeypatch.setattr(runner_module, "AsyncSessionLocal", lambda: _FakeSessionCtx())
monkeypatch.setattr(
runner, "_build_memory_stage_config", _fake_build_memory_stage_config
)
monkeypatch.setattr(runner, "_build_stage_toolkit", lambda **kwargs: object())
monkeypatch.setattr(
runner,
"_execute_single_stage_step",
_fake_execute_single_stage_step,
)
result = await runner.execute(
user_context=_user_context(),
context_messages=[],
pipeline=_FakePipeline(),
run_input=_run_input(),
system_agent_mode="memory",
memory_job_config=default_memory_job_config(),
)
assert result["memory"]["answer"] == "memory"
@@ -18,7 +18,7 @@ def _run_input_payload() -> dict[str, Any]:
"messages": [],
"tools": [],
"context": [],
"forwardedProps": {},
"forwardedProps": {"agent_type": "worker"},
}
@@ -35,7 +35,7 @@ async def _fake_user_context(**kwargs: object) -> UserContext:
return UserContext(
id=str(uuid4()),
username="alice",
email="alice@example.com",
phone="+8613900000000",
settings=parse_profile_settings(None),
)
@@ -177,17 +177,53 @@ async def test_run_agentscope_task_rejects_invalid_command_type() -> None:
)
@pytest.mark.asyncio
async def test_run_agentscope_task_requires_forwarded_props_agent_type() -> None:
payload = _run_input_payload()
payload["forwardedProps"] = {}
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
await tasks_module.run_agentscope_task(
{
"command": "run",
"owner_id": str(uuid4()),
"run_input": payload,
}
)
@pytest.mark.asyncio
async def test_run_agentscope_task_memory_mode_requires_automation_job_id() -> None:
payload = _run_input_payload()
payload["forwardedProps"] = {"agent_type": "memory"}
with pytest.raises(
ValueError, match="automation_job_id is required for memory mode"
):
await tasks_module.run_agentscope_task(
{
"command": "run",
"owner_id": str(uuid4()),
"run_input": payload,
}
)
@pytest.mark.asyncio
async def test_build_recent_context_messages_includes_all_user_attachments(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakeAgentService:
async def load_agent_input_messages(
class _FakeContextService:
def __init__(self, *, repository: object) -> None:
del repository
async def load_context_messages(
self,
*,
thread_id: str,
system_agent_mode: str,
) -> dict[str, object] | None:
del thread_id
del thread_id, system_agent_mode
return {
"messages": [
{
@@ -215,14 +251,13 @@ async def test_build_recent_context_messages_includes_all_user_attachments(
async def download_bytes(self, *, bucket: str, path: str) -> bytes:
return f"{bucket}:{path}".encode("utf-8")
monkeypatch.setattr(
tasks_module, "get_agent_service", lambda session: _FakeAgentService()
)
monkeypatch.setattr(tasks_module, "AgentContextService", _FakeContextService)
monkeypatch.setattr(tasks_module, "supabase_service", _FakeSupabase())
messages = await tasks_module._build_recent_context_messages(
session=object(),
thread_id=str(uuid4()),
context_mode="worker",
)
assert len(messages) == 1
@@ -238,13 +273,17 @@ async def test_build_recent_context_messages_includes_all_user_attachments(
async def test_build_recent_context_messages_uses_tool_metadata_output(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakeAgentService:
async def load_agent_input_messages(
class _FakeContextService:
def __init__(self, *, repository: object) -> None:
del repository
async def load_context_messages(
self,
*,
thread_id: str,
system_agent_mode: str,
) -> dict[str, object] | None:
del thread_id
del thread_id, system_agent_mode
return {
"messages": [
{
@@ -268,13 +307,12 @@ async def test_build_recent_context_messages_uses_tool_metadata_output(
]
}
monkeypatch.setattr(
tasks_module, "get_agent_service", lambda session: _FakeAgentService()
)
monkeypatch.setattr(tasks_module, "AgentContextService", _FakeContextService)
messages = await tasks_module._build_recent_context_messages(
session=object(),
thread_id=str(uuid4()),
context_mode="worker",
)
assert len(messages) == 1
@@ -290,13 +328,17 @@ async def test_build_recent_context_messages_uses_tool_metadata_output(
async def test_build_recent_context_messages_skips_tool_without_metadata_output(
monkeypatch: pytest.MonkeyPatch,
) -> None:
class _FakeAgentService:
async def load_agent_input_messages(
class _FakeContextService:
def __init__(self, *, repository: object) -> None:
del repository
async def load_context_messages(
self,
*,
thread_id: str,
system_agent_mode: str,
) -> dict[str, object] | None:
del thread_id
del thread_id, system_agent_mode
return {
"messages": [
{
@@ -307,13 +349,44 @@ async def test_build_recent_context_messages_skips_tool_without_metadata_output(
]
}
monkeypatch.setattr(
tasks_module, "get_agent_service", lambda session: _FakeAgentService()
)
monkeypatch.setattr(tasks_module, "AgentContextService", _FakeContextService)
messages = await tasks_module._build_recent_context_messages(
session=object(),
thread_id=str(uuid4()),
context_mode="worker",
)
assert messages == []
@pytest.mark.asyncio
async def test_build_recent_context_messages_passes_context_mode_through(
monkeypatch: pytest.MonkeyPatch,
) -> None:
captured_mode: dict[str, str | None] = {"mode": None}
class _FakeContextService:
def __init__(self, *, repository: object) -> None:
del repository
async def load_context_messages(
self,
*,
thread_id: str,
system_agent_mode: str,
) -> dict[str, object] | None:
del thread_id
captured_mode["mode"] = system_agent_mode
return None
monkeypatch.setattr(tasks_module, "AgentContextService", _FakeContextService)
messages = await tasks_module._build_recent_context_messages(
session=object(),
thread_id=str(uuid4()),
context_mode="worker",
)
assert messages == []
assert captured_mode["mode"] == "worker"
@@ -1,99 +1,8 @@
from __future__ import annotations
import pytest
from pydantic import ValidationError
from core.agentscope import schemas as exported_schemas
from core.agentscope.schemas.agent_runtime import (
AcceptedTaskResponse,
AgUiWireEvent,
HistorySnapshot,
HistorySnapshotResponse,
InternalRuntimeEvent,
RunCommand,
pytest.skip(
"legacy agent_runtime schemas removed; covered by agui_input tests",
allow_module_level=True,
)
def test_run_command_alias_roundtrip() -> None:
payload = {
"threadId": "thread-001",
"runId": "run-001",
"state": {"cursor": 1},
"messages": [{"role": "user", "content": "hi"}],
"tools": [{"name": "calendar.lookup"}],
"context": {"locale": "zh-CN"},
"forwardedProps": {"traceId": "trace-1"},
}
command = RunCommand.model_validate(payload)
assert command.thread_id == "thread-001"
assert command.run_id == "run-001"
assert command.forwarded_props == {"traceId": "trace-1"}
dumped = command.model_dump(mode="json", by_alias=True)
assert dumped["threadId"] == "thread-001"
assert dumped["runId"] == "run-001"
assert dumped["forwardedProps"] == {"traceId": "trace-1"}
def test_history_snapshot_response_shape() -> None:
response = HistorySnapshotResponse(
threadId="thread-123",
snapshot=HistorySnapshot(
threadId="thread-123",
day="2026-03-11",
hasMore=False,
messages=[{"id": "msg-1"}],
),
)
dumped = response.model_dump(mode="json", by_alias=True, exclude_none=True)
assert dumped["type"] == "STATE_SNAPSHOT"
assert dumped["threadId"] == "thread-123"
assert dumped["snapshot"]["scope"] == "history_day"
assert dumped["snapshot"]["hasMore"] is False
assert dumped["snapshot"]["messages"] == [{"id": "msg-1"}]
def test_runtime_event_validation_basics() -> None:
internal = InternalRuntimeEvent(type="RUN_STARTED", data={"step": 1})
assert internal.type == "RUN_STARTED"
assert internal.model_dump(mode="json", by_alias=True)["data"] == {"step": 1}
wire = AgUiWireEvent(type="TEXT_MESSAGE_CONTENT", payload={"delta": "hello"})
dumped = wire.model_dump(mode="json", by_alias=True, exclude_none=True)
assert dumped["type"] == "TEXT_MESSAGE_CONTENT"
assert dumped["payload"] == {"delta": "hello"}
with pytest.raises(ValidationError):
InternalRuntimeEvent.model_validate({"threadId": "t-1", "data": {}})
with pytest.raises(ValidationError):
AgUiWireEvent.model_validate({"payload": {"delta": "hello"}})
def test_schemas_exports_include_task_and_history_models() -> None:
assert exported_schemas.AcceptedTaskResponse is AcceptedTaskResponse
assert exported_schemas.TaskAccepted is AcceptedTaskResponse
assert exported_schemas.TaskAcceptedResponse is AcceptedTaskResponse
assert exported_schemas.HistorySnapshotResponse is HistorySnapshotResponse
def test_run_command_accepts_agui_context_list_and_parent_run_id() -> None:
payload = {
"threadId": "thread-xyz",
"runId": "run-xyz",
"state": {},
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
"tools": [],
"context": [],
"forwardedProps": {},
"parentRunId": None,
}
command = RunCommand.model_validate(payload)
dumped = command.model_dump(mode="json", by_alias=True)
assert dumped["context"] == []
assert "parentRunId" in dumped
@@ -20,7 +20,7 @@ def _base_payload() -> dict[str, object]:
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
"tools": [],
"context": [],
"forwardedProps": {},
"forwardedProps": {"agent_type": "worker"},
}
@@ -149,7 +149,7 @@ def test_parse_run_input_accepts_snake_case_aliases() -> None:
],
"tools": [],
"context": [],
"forwarded_props": {},
"forwarded_props": {"agent_type": "worker"},
}
run_input = parse_run_input(payload)
@@ -162,11 +162,12 @@ def test_parse_run_input_accepts_snake_case_aliases() -> None:
def test_parse_run_input_accepts_client_time_forwarded_props() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"agent_type": "worker",
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16T09:12:33-07:00",
"client_epoch_ms": 1773658353000,
}
},
}
run_input = parse_run_input(payload)
@@ -177,11 +178,12 @@ def test_parse_run_input_accepts_client_time_forwarded_props() -> None:
def test_parse_run_input_rejects_invalid_client_time_timezone() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"agent_type": "worker",
"client_time": {
"device_timezone": "Mars/OlympusMons",
"client_now_iso": "2026-03-16T09:12:33-07:00",
"client_epoch_ms": 1773658353000,
}
},
}
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
@@ -191,11 +193,12 @@ def test_parse_run_input_rejects_invalid_client_time_timezone() -> None:
def test_parse_run_input_rejects_invalid_client_time_now_iso() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"agent_type": "worker",
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16 09:12:33",
"client_epoch_ms": 1773658353000,
}
},
}
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
@@ -205,11 +208,12 @@ def test_parse_run_input_rejects_invalid_client_time_now_iso() -> None:
def test_parse_run_input_rejects_invalid_client_time_epoch_type() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"agent_type": "worker",
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16T09:12:33-07:00",
"client_epoch_ms": "1773658353000",
}
},
}
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
@@ -219,6 +223,7 @@ def test_parse_run_input_rejects_invalid_client_time_epoch_type() -> None:
def test_parse_run_input_rejects_unknown_forwarded_props_key() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"agent_type": "worker",
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16T09:12:33-07:00",
@@ -229,3 +234,17 @@ def test_parse_run_input_rejects_unknown_forwarded_props_key() -> None:
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
parse_run_input(payload)
def test_parse_run_input_rejects_missing_forwarded_props_agent_type() -> None:
payload = _base_payload()
payload["forwardedProps"] = {
"client_time": {
"device_timezone": "America/Los_Angeles",
"client_now_iso": "2026-03-16T09:12:33-07:00",
"client_epoch_ms": 1773658353000,
}
}
with pytest.raises(ValueError, match="invalid RunAgentInput.forwardedProps"):
parse_run_input(payload)
@@ -11,7 +11,7 @@ def test_build_agent_prompt_for_worker_contains_runtime_config() -> None:
{
"temperature": 0.2,
"context_messages": {"mode": "number", "count": 20},
"enabled_tool_groups": ["read", "write"],
"enabled_tools": ["calendar.read", "calendar.write"],
}
),
)
@@ -20,7 +20,7 @@ def test_build_agent_prompt_for_worker_contains_runtime_config() -> None:
assert "- type: worker" in prompt
assert "context_messages.mode=number" in prompt
assert "context_messages.count=20" in prompt
assert "enabled_tool_groups=read,write" in prompt
assert "enabled_tools=calendar.read,calendar.write" in prompt
def test_build_agent_prompt_for_memory_uses_memory_rules() -> None:
@@ -29,7 +29,7 @@ def test_build_agent_prompt_for_memory_uses_memory_rules() -> None:
llm_config=SystemAgentLLMConfig.model_validate(
{
"context_messages": {"mode": "day", "count": 2},
"enabled_tool_groups": ["read"],
"enabled_tools": ["user.lookup"],
}
),
)
@@ -38,4 +38,4 @@ def test_build_agent_prompt_for_memory_uses_memory_rules() -> None:
assert "[Memory Agent]" in prompt
assert "context_messages.mode=day" in prompt
assert "context_messages.count=2" in prompt
assert "enabled_tool_groups=read" in prompt
assert "enabled_tools=user.lookup" in prompt
@@ -1,11 +1,12 @@
from __future__ import annotations
import json
from typing import Any, AsyncGenerator
import pytest
from core.agentscope.tools.hitl_middleware import create_hitl_middleware
from core.agentscope.tools.tool_meta import TOOL_META, ToolMeta
from core.agentscope.tools.tool_config import ToolApprovalConfig, ToolConfig, ToolGroup
from core.agentscope.tools.tool_middleware import create_approval_middleware
async def _next_handler(**kwargs: Any) -> AsyncGenerator[dict[str, object], None]:
@@ -15,9 +16,31 @@ async def _next_handler(**kwargs: Any) -> AsyncGenerator[dict[str, object], None
return _generator()
def _extract_error_payload(chunk: object) -> dict[str, Any]:
content = getattr(chunk, "content", None)
if not isinstance(content, list) or not content:
return {}
first_block = content[0]
text = getattr(first_block, "text", None)
if not isinstance(text, str) and isinstance(first_block, dict):
raw_text = first_block.get("text")
text = raw_text if isinstance(raw_text, str) else None
if not isinstance(text, str):
return {}
return json.loads(text)
@pytest.mark.asyncio
async def test_hitl_middleware_default_write_does_not_require_approval() -> None:
middleware = create_hitl_middleware(meta_by_name=TOOL_META)
middleware = create_approval_middleware(
config_by_name={
"calendar_write": ToolConfig(
name="calendar_write",
group=ToolGroup.EXECUTE,
approval=ToolApprovalConfig(required=False),
)
}
)
responses = []
async for chunk in middleware(
@@ -30,36 +53,39 @@ async def test_hitl_middleware_default_write_does_not_require_approval() -> None
@pytest.mark.asyncio
async def test_hitl_middleware_pending_when_tool_requires_approval(
monkeypatch: pytest.MonkeyPatch,
) -> None:
middleware = create_hitl_middleware(
meta_by_name={
"calendar.write": ToolMeta(name="calendar.write", requires_approval=True)
async def test_hitl_middleware_pending_when_tool_requires_approval() -> None:
middleware = create_approval_middleware(
config_by_name={
"calendar_write": ToolConfig(
name="calendar_write",
group=ToolGroup.EXECUTE,
approval=ToolApprovalConfig(required=True),
)
}
)
monkeypatch.setattr(
"core.agentscope.tools.hitl_middleware.build_tool_response",
lambda payload: payload,
)
responses = []
async for chunk in middleware(
{"tool_call": {"name": "calendar.write", "input": {"operation": "create"}}},
{"tool_call": {"name": "calendar_write", "input": {"operation": "create"}}},
_next_handler,
):
responses.append(chunk)
assert responses[0]["data"]["status"] == "pending"
payload = _extract_error_payload(responses[0])
assert payload["error"]["code"] == "TOOL_PENDING_APPROVAL"
@pytest.mark.asyncio
async def test_hitl_middleware_passes_when_write_approved() -> None:
middleware = create_hitl_middleware(
meta_by_name={
"calendar.write": ToolMeta(name="calendar.write", requires_approval=True)
middleware = create_approval_middleware(
config_by_name={
"calendar_write": ToolConfig(
name="calendar_write",
group=ToolGroup.EXECUTE,
approval=ToolApprovalConfig(required=True),
)
},
approval_resolver=lambda _name, _args: "approved",
approval_resolver=lambda _name, _args, _config: "approved",
)
responses = []
@@ -69,6 +95,7 @@ async def test_hitl_middleware_passes_when_write_approved() -> None:
"name": "calendar.write",
"input": {
"operation": "create",
"_hitl": {"approval": "required"},
},
}
},
@@ -82,25 +109,24 @@ async def test_hitl_middleware_passes_when_write_approved() -> None:
@pytest.mark.asyncio
async def test_hitl_middleware_rejected_short_circuits(
monkeypatch: pytest.MonkeyPatch,
) -> None:
middleware = create_hitl_middleware(
meta_by_name={
"calendar.write": ToolMeta(name="calendar.write", requires_approval=True)
async def test_hitl_middleware_rejected_short_circuits() -> None:
middleware = create_approval_middleware(
config_by_name={
"calendar_write": ToolConfig(
name="calendar_write",
group=ToolGroup.EXECUTE,
approval=ToolApprovalConfig(required=True),
)
},
approval_resolver=lambda _name, _args: "rejected",
)
monkeypatch.setattr(
"core.agentscope.tools.hitl_middleware.build_tool_response",
lambda payload: payload,
approval_resolver=lambda _name, _args, _config: "rejected",
)
responses = []
async for chunk in middleware(
{"tool_call": {"name": "calendar.write", "input": {"operation": "create"}}},
{"tool_call": {"name": "calendar_write", "input": {"operation": "create"}}},
_next_handler,
):
responses.append(chunk)
assert responses[0]["data"]["status"] == "rejected"
payload = _extract_error_payload(responses[0])
assert payload["error"]["code"] == "TOOL_REJECTED"
@@ -16,7 +16,7 @@ def _build_user_context(*, timezone_name: str = "Asia/Shanghai") -> UserContext:
return UserContext(
id=str(uuid4()),
username="alice",
email="alice@example.com",
phone="+8613900000000",
bio="focus on calendars",
settings=parse_profile_settings(
{
@@ -7,7 +7,9 @@ from core.agentscope.tools.toolkit import build_stage_toolkit
from schemas.agent.system_agent import AgentType
def test_build_stage_toolkit_filters_requested_tools_by_agent_type(monkeypatch) -> None:
def test_build_stage_toolkit_uses_explicit_enabled_tools_as_final_set(
monkeypatch,
) -> None:
captured: dict[str, object] = {}
def _fake_build_toolkit(**kwargs):
@@ -22,7 +24,29 @@ def test_build_stage_toolkit_filters_requested_tools_by_agent_type(monkeypatch)
agent_type=AgentType.WORKER,
session=cast(Any, object()),
owner_id=uuid4(),
enabled_tool_names={"calendar_read", "calendar_write", "user_lookup"},
enabled_tool_names={"calendar_read", "user_lookup"},
)
assert captured["enabled_tool_names"] == {"calendar_read", "user_lookup"}
def test_build_stage_toolkit_uses_memory_defaults_without_explicit_tools(
monkeypatch,
) -> None:
captured: dict[str, object] = {}
def _fake_build_toolkit(**kwargs):
captured.update(kwargs)
return object()
monkeypatch.setattr(
"core.agentscope.tools.toolkit.build_toolkit", _fake_build_toolkit
)
build_stage_toolkit(
agent_type=AgentType.MEMORY,
session=cast(Any, object()),
owner_id=uuid4(),
)
assert captured["enabled_tool_names"] == {"calendar_read", "user_lookup"}
@@ -0,0 +1,134 @@
from __future__ import annotations
from datetime import datetime, timedelta, timezone
from uuid import UUID, uuid4
import pytest
from core.automation.scheduler import (
AutomationSchedulerService,
_compute_next_run_at,
)
from models.automation_jobs import ScheduleType
from schemas.automation.config import AutomationJobConfig
from schemas.automation.scheduler import DueAutomationJob
class _FakeRepository:
def __init__(self, jobs: list[DueAutomationJob]) -> None:
self.jobs = jobs
self.marked: list[tuple[UUID, datetime, datetime]] = []
self.commits = 0
self.rollbacks = 0
async def list_due_jobs(
self, *, now_utc: datetime, limit: int
) -> list[DueAutomationJob]:
del now_utc
return self.jobs[:limit]
async def get_job_config(self, *, job_id: UUID) -> AutomationJobConfig:
del job_id
return AutomationJobConfig.model_validate(
{
"agent_type": "memory",
"model_code": "qwen3.5-flash",
"enabled_tools": ["calendar.read", "user.lookup"],
"input_template": "auto input",
"context": {
"source": "latest_chat",
"window_mode": "day",
"window_count": 2,
},
}
)
async def ensure_latest_chat_session(self, *, owner_id: UUID) -> UUID:
return owner_id
async def mark_job_dispatched(
self,
*,
job_id: UUID,
next_run_at: datetime,
last_run_at: datetime,
) -> None:
self.marked.append((job_id, next_run_at, last_run_at))
async def commit(self) -> None:
self.commits += 1
async def rollback(self) -> None:
self.rollbacks += 1
class _FakeQueue:
def __init__(self) -> None:
self.commands: list[dict[str, object]] = []
async def enqueue(
self,
*,
command: dict[str, object],
dedup_key: str | None,
) -> str:
del dedup_key
self.commands.append(command)
return "task-1"
@pytest.mark.asyncio
async def test_scan_and_dispatch_enqueues_memory_run_command() -> None:
now = datetime(2026, 3, 19, 12, 0, tzinfo=timezone.utc)
owner_id = uuid4()
job_id = uuid4()
repo = _FakeRepository(
jobs=[
DueAutomationJob(
id=job_id,
owner_id=owner_id,
schedule_type=ScheduleType.DAILY,
timezone="UTC",
next_run_at=now - timedelta(minutes=1),
)
]
)
queue = _FakeQueue()
service = AutomationSchedulerService(repository=repo, queue=queue)
result = await service.scan_and_dispatch(now_utc=now, limit=10)
assert result.scanned == 1
assert result.dispatched == 1
assert len(queue.commands) == 1
run_input = queue.commands[0]["run_input"]
assert isinstance(run_input, dict)
assert run_input["forwardedProps"] == {"agent_type": "memory"}
assert queue.commands[0]["automation_job_id"] == str(job_id)
assert repo.commits == 1
def test_compute_next_run_at_daily() -> None:
now = datetime(2026, 3, 19, 12, 0, tzinfo=timezone.utc)
current = datetime(2026, 3, 19, 11, 0, tzinfo=timezone.utc)
computed = _compute_next_run_at(
current_next_run_at=current,
now_utc=now,
schedule_type=ScheduleType.DAILY,
)
assert computed == datetime(2026, 3, 20, 11, 0, tzinfo=timezone.utc)
def test_compute_next_run_at_weekly() -> None:
now = datetime(2026, 3, 19, 12, 0, tzinfo=timezone.utc)
current = datetime(2026, 3, 10, 11, 0, tzinfo=timezone.utc)
computed = _compute_next_run_at(
current_next_run_at=current,
now_utc=now,
schedule_type=ScheduleType.WEEKLY,
)
assert computed == datetime(2026, 3, 24, 11, 0, tzinfo=timezone.utc)
@@ -0,0 +1,18 @@
from __future__ import annotations
from pathlib import Path
def test_memory_automation_job_trigger_exists_in_0004_migration() -> None:
migration = (
Path(__file__).resolve().parents[3]
/ "alembic"
/ "versions"
/ "20260319_0004_automation_job_config_for_memory.py"
)
content = migration.read_text(encoding="utf-8")
assert "INSERT INTO public.automation_jobs" in content
assert "'agent_type', 'memory'" in content
assert "ux_automation_jobs_owner_memory_active" in content
assert "input_template" in content
@@ -0,0 +1,37 @@
from __future__ import annotations
import pytest
from schemas.agent.consumer_registry import AgentConsumerBinding, ConsumerRegistry
from schemas.agent.pipeline_spec import (
ContextPolicy,
ExecutorKind,
PipelineSpec,
StageSpec,
)
def test_consumer_registry_rejects_duplicate_bits() -> None:
with pytest.raises(ValueError, match="duplicate visibility bit"):
ConsumerRegistry(
bindings=[
AgentConsumerBinding(agent_type="router", bit=16),
AgentConsumerBinding(agent_type="worker", bit=16),
]
)
def test_pipeline_spec_requires_non_empty_stages() -> None:
with pytest.raises(ValueError, match="at least 1 item"):
PipelineSpec(mode="worker", stages=[])
def test_stage_spec_normalizes_stage_name() -> None:
spec = StageSpec(
stage_name=" Worker ",
executor_kind=ExecutorKind.REACT,
default_visibility_mask=1,
context_policy=ContextPolicy(consumer_agent_type="worker", count=20),
)
assert spec.stage_name == "worker"
@@ -0,0 +1,31 @@
from __future__ import annotations
import pytest
from schemas.agent.system_agent import SystemAgentLLMConfig
def test_system_agent_llm_config_normalizes_enabled_tools_aliases() -> None:
config = SystemAgentLLMConfig.model_validate(
{
"enabled_tools": [
"calendar.write",
"calendar_write",
"user.lookup",
]
}
)
assert [tool.value for tool in config.enabled_tools] == [
"calendar.write",
"user.lookup",
]
def test_system_agent_llm_config_rejects_unknown_enabled_tool() -> None:
with pytest.raises(ValueError, match="unknown enabled tool"):
SystemAgentLLMConfig.model_validate(
{
"enabled_tools": ["calendar.remove"],
}
)
@@ -0,0 +1,23 @@
from __future__ import annotations
import pytest
from schemas.agent.visibility import VisibilityMask, bit_mask
def test_visibility_mask_from_bits_and_contains() -> None:
mask = VisibilityMask.from_bits(bits=[0, 16, 18])
assert mask.contains(bit=0) is True
assert mask.contains(bit=16) is True
assert mask.contains(bit=17) is False
def test_visibility_mask_rejects_out_of_range_bit() -> None:
with pytest.raises(ValueError, match="range"):
VisibilityMask.from_bits(bits=[64])
def test_bit_mask_builds_single_bit_integer() -> None:
assert bit_mask(bit=0) == 1
assert bit_mask(bit=16) == (1 << 16)
@@ -0,0 +1,30 @@
from __future__ import annotations
import pytest
from schemas.automation.config import AutomationJobConfig, default_memory_job_config
def test_default_memory_job_config_has_expected_defaults() -> None:
config = default_memory_job_config()
assert config.agent_type.value == "memory"
assert config.model_code == "qwen3.5-flash"
assert config.context.source.value == "latest_chat"
def test_automation_job_config_rejects_non_flash_model() -> None:
with pytest.raises(ValueError, match="model_code must be qwen3.5-flash"):
AutomationJobConfig.model_validate(
{
"agent_type": "memory",
"model_code": "qwen-plus",
"enabled_tools": ["calendar.read"],
"input_template": "x",
"context": {
"source": "latest_chat",
"window_mode": "day",
"window_count": 2,
},
}
)