by-framework 0.2.2.dev9__tar.gz → 0.2.2.dev10__tar.gz

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: by-framework
-Version: 0.2.2.dev9
+Version: 0.2.2.dev10
 Summary: 分布式 Agent 调度框架
 License-File: LICENSE
 Requires-Python: >=3.12

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "by-framework"
-version = "0.2.2.dev9"
+version = "0.2.2.dev10"
 description = "分布式 Agent 调度框架"
 readme = "README.md"
 requires-python = ">=3.12"

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/__init__.py

@@ -8,6 +8,7 @@ from `GatewayWorker` and running `run_worker`.
 from .admin import WorkerManager
 from .client.byai_client import ByaiGatewayClient
 from .client.client import (
+    CancelSessionResponse,
     CancelTaskResponse,
     DataStreamEntry,
     GatewayClient,
@@ -148,6 +149,7 @@ __all__ = [
     "DataStreamEntry",
     "SendMessageResponse",
     "CancelTaskResponse",
+    "CancelSessionResponse",
     "run_worker",
     "logger",
     "setup_logging",

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/client/client.py

@@ -37,6 +37,7 @@ from by_framework.core.protocol.commands import (
 from by_framework.core.protocol.data_message import DataMessage
 from by_framework.core.protocol.message_header import MessageHeader
 from by_framework.core.protocol.responses import (
+    CancelSessionResponse,
     CancelTaskResponse,
     ExecutionStatus,
     SendMessageResponse,
@@ -581,6 +582,99 @@ class GatewayClient:
             cancelled_count=len(to_cancel),
         )
 
+    async def cancel_session(
+        self,
+        session_id: str,
+        reason: str = "",
+        requested_by: str = "client",
+        cancel_mode: str = CancelMode.GRACEFUL,
+    ) -> CancelSessionResponse:
+        """Cancel every active execution registered under a session.
+
+        Unlike cancel_task, this does not require a starting message_id:
+        it cancels the whole session's active executions in one call.
+        """
+        if self.registry is None:
+            raise ValueError("GatewayClient requires a WorkerRegistry to cancel tasks")
+
+        all_executions = await self.registry.get_all_session_executions(session_id)
+
+        if not all_executions:
+            return CancelSessionResponse(
+                success=False,
+                session_id=session_id,
+                status=ExecutionStatus.NOT_FOUND,
+                timestamp=int(time.time() * 1000),
+                cancelled_count=0,
+                already_finished_count=0,
+            )
+
+        terminal_states = {"COMPLETED", "FAILED", "CANCELLED"}
+        to_cancel = [
+            node
+            for node in all_executions
+            if node.get("status", "") not in terminal_states
+        ]
+        terminal_nodes = [
+            node for node in all_executions if node.get("status", "") in terminal_states
+        ]
+
+        # Flag terminal executions so a still-running descendant can't wake
+        # them back up via callback, without changing their status.
+        for node in terminal_nodes:
+            await self.registry.mark_cancel_requested(
+                node.get("execution_id", ""), session_id, reason
+            )
+
+        if not to_cancel:
+            return CancelSessionResponse(
+                success=False,
+                session_id=session_id,
+                status=ExecutionStatus.ALREADY_FINISHED,
+                timestamp=int(time.time() * 1000),
+                cancelled_count=0,
+                already_finished_count=len(terminal_nodes),
+            )
+
+        for node in to_cancel:
+            node_execution_id = node.get("execution_id", "")
+            node_worker_id = node.get("worker_id", "")
+            node_message_id = node.get("message_id", "")
+
+            await self.registry.mark_execution_cancelling(
+                node_execution_id, session_id, reason
+            )
+
+            if node_worker_id:
+                cancel_command = CancelTaskCommand(
+                    header=MessageHeader(
+                        message_id=f"{CANCEL_MESSAGE_ID_PREFIX}{uuid.uuid4().hex[:8]}",
+                        session_id=session_id,
+                        trace_id=node.get("trace_id") or uuid.uuid4().hex,
+                        target_agent_type=node.get("target_agent_type", ""),
+                        parent_message_id=node_message_id,
+                    ),
+                    target_message_id=node_message_id,
+                    target_execution_id=node_execution_id,
+                    target_worker_id=node_worker_id,
+                    reason=reason,
+                    requested_by=requested_by,
+                    cancel_mode=cancel_mode,
+                )
+                await self.redis.xadd(
+                    RedisKeys.worker_ctrl_stream(node_worker_id),
+                    cancel_command.to_redis_payload(),
+                )
+
+        return CancelSessionResponse(
+            success=True,
+            session_id=session_id,
+            status=ExecutionStatus.CANCEL_REQUESTED,
+            timestamp=int(time.time() * 1000),
+            cancelled_count=len(to_cancel),
+            already_finished_count=len(terminal_nodes),
+        )
+
     async def send_message(
         self,
         target_agent_type: str,
@@ -691,7 +785,28 @@ class GatewayClient:
             content=params["content"],
             extra_payload=params["extra_payload"],
         )
-        execution_id = f"{EXECUTION_ID_PREFIX}{uuid.uuid4().hex[:8]}"
+
+        # A RESUME reuses the message_id of the original AskAgentCommand so the
+        # worker can look the suspended execution back up. Reuse its
+        # execution_id too, and skip re-initializing the registry record for
+        # it below -- initialize_execution() would otherwise overwrite the
+        # message_id -> execution_id mapping and detach this resume from the
+        # execution it's meant to continue.
+        resumed_execution = None
+        if (
+            params["action_type"] == ActionType.RESUME.value
+            and self.registry
+            and hasattr(self.registry, "get_execution_by_message_id")
+        ):
+            resumed_execution = await self.registry.get_execution_by_message_id(
+                message_id, session_id=params["session_id"]
+            )
+
+        execution_id = (
+            resumed_execution["execution_id"]
+            if resumed_execution
+            else f"{EXECUTION_ID_PREFIX}{uuid.uuid4().hex[:8]}"
+        )
 
         # 3. Resolve route and optionally probe agent type/liveness
         should_dispatch_control = True
@@ -837,8 +952,14 @@ class GatewayClient:
                 error_code=ExecutionStatus.ERR_AGENT_TYPE_UNAVAILABLE,
             )
 
-        # Initialize execution tracking
-        if self.registry and hasattr(self.registry, "initialize_execution"):
+        # Initialize execution tracking. Skipped for a RESUME that reattached
+        # to an existing execution -- that execution is already tracked, and
+        # re-initializing it here would overwrite its message_id mapping.
+        if (
+            not resumed_execution
+            and self.registry
+            and hasattr(self.registry, "initialize_execution")
+        ):
             try:
                 await self.registry.initialize_execution(
                     {

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/common/emitter.py

@@ -107,6 +107,22 @@ class DefaultSseLayoutBuilder(DataLayoutBuilder):
         parent_order_id: Optional[str] = None,
     ) -> Dict[str, Any]:
         """Build default BaiYing-compatible user input form payload."""
+        try:
+            parsed_prompt = json.loads(prompt)
+        except (json.JSONDecodeError, TypeError):
+            parsed_prompt = None
+
+        # Structured question prompts are already in the format expected by clients.
+        if isinstance(parsed_prompt, dict) and "questions" in parsed_prompt:
+            return self.build(
+                content=prompt,
+                role="assistant",
+                content_type=SseReasonMessageType.ask_user_question.value,
+                source_agent_type=source_agent_type,
+                order_id=order_id,
+                parent_order_id=parent_order_id,
+            )
+
         input_form = {
             "formStatus": 0,
             "pluginMachineFields": [

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/core/__init__.py

@@ -26,6 +26,7 @@ from .protocol import (
     BaiYingMessage,
     BaiYingMessageRole,
     BaseCommand,
+    CancelSessionResponse,
     CancelTaskCommand,
     CancelTaskResponse,
     DataMessage,
@@ -53,6 +54,7 @@ from .workspace import WorkspaceManager
 __all__ = [
     "SendMessageResponse",
     "CancelTaskResponse",
+    "CancelSessionResponse",
     "BaiYingMessage",
     "BaiYingMessageRole",
     "MessageContent",

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/core/protocol/__init__.py

@@ -65,6 +65,7 @@ from .message import (
 )
 from .message_header import MessageHeader
 from .responses import (
+    CancelSessionResponse,
     CancelTaskResponse,
     CancelTaskResponseDict,
     SendMessageResponse,
@@ -81,6 +82,7 @@ from .results import (
 __all__ = [
     "SendMessageResponse",
     "CancelTaskResponse",
+    "CancelSessionResponse",
     "SendMessageResponseDict",
     "CancelTaskResponseDict",
     "AgentTaskResult",

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/core/protocol/content_type.py

@@ -33,6 +33,7 @@ class SseReasonMessageType(str, Enum):
     think_code_result = "3007"  # thinking process code execution result
     task_finished = "3009"  # task finished
     task_user_input = "3013"  # user input
+    ask_user_question = "3014"  # ask user question
     task_create_file = "3010"  # create file
     task_title = "3011"  # task title
     agent_card = "2015"  # agent card

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/core/protocol/responses.py

@@ -92,3 +92,14 @@ class CancelTaskResponse:
     timestamp: int
     error: str = ""
     cancelled_count: int = 0
+
+
+@dataclass(frozen=True)
+class CancelSessionResponse:
+    success: bool
+    session_id: str
+    status: str
+    timestamp: int
+    cancelled_count: int = 0
+    already_finished_count: int = 0
+    error: str = ""

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/core/registry.py

@@ -1011,7 +1011,8 @@ class WorkerRegistry:
         old_execution = dict(current)
         current["status"] = status
         now = int(time.time() * 1000)
-        current["finished_at"] = now
+        if is_terminal_state(status):
+            current["finished_at"] = now
         current["updated_at"] = now
         if completion:
             for key, value in completion.items():

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/src/by_framework/worker/runner.py

@@ -505,6 +505,16 @@ class WorkerRunner:
                     header.message_id, session_id=header.session_id
                 )
 
+            if isinstance(command, ResumeCommand) and existing_execution is None:
+                logger.warning(
+                    "[%s] ResumeCommand did not resolve to an existing execution "
+                    "(message_id=%s, session_id=%s); starting a new, disconnected "
+                    "execution instead of continuing the suspended one.",
+                    self.worker.worker_id,
+                    header.message_id,
+                    header.session_id,
+                )
+
             # Skip terminal state replays
             if (
                 existing_execution

{by_framework-0.2.2.dev9 → by_framework-0.2.2.dev10}/tests/client/test_client.py

@@ -1022,3 +1022,348 @@ async def test_client_cascading_cancel_root_completed_children_running():
     registry.mark_cancel_requested.assert_called_once_with(
         "exec-a", "sess-1", "user abort"
     )
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_cancels_all_non_terminal_executions():
+    """Test that cancel_session cancels every non-terminal execution in a
+    session in one call, without needing a starting message_id."""
+    exec_a = {
+        "execution_id": "exec-a",
+        "message_id": "msg-a",
+        "session_id": "sess-1",
+        "worker_id": "worker-1",
+        "target_agent_type": "agent_a",
+        "status": "RUNNING",
+        "parent_message_id": "",
+    }
+    exec_b = {
+        "execution_id": "exec-b",
+        "message_id": "msg-b",
+        "session_id": "sess-1",
+        "worker_id": "worker-2",
+        "target_agent_type": "agent_b",
+        "status": "QUEUED",
+        "parent_message_id": "msg-a",
+    }
+
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = [exec_a, exec_b]
+    redis = AsyncMock()
+    client = GatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-1", reason="user abort")
+
+    assert result.success is True
+    assert result.status == "CANCEL_REQUESTED"
+    assert result.cancelled_count == 2
+    assert result.already_finished_count == 0
+
+    assert registry.mark_execution_cancelling.call_count == 2
+    cancel_ids = [
+        call.args[0] for call in registry.mark_execution_cancelling.call_args_list
+    ]
+    assert "exec-a" in cancel_ids
+    assert "exec-b" in cancel_ids
+
+    assert redis.xadd.call_count == 2
+    target_streams = [call.args[0] for call in redis.xadd.call_args_list]
+    assert any("worker-1" in s for s in target_streams)
+    assert any("worker-2" in s for s in target_streams)
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_flags_terminal_executions_without_cancelling():
+    """Terminal executions in the session should only be flagged with
+    cancel_requested (not re-marked cancelling, no CancelTaskCommand sent),
+    so a still-running descendant can't wake them back up via callback."""
+    exec_a = {
+        "execution_id": "exec-a",
+        "message_id": "msg-a",
+        "session_id": "sess-1",
+        "worker_id": "worker-1",
+        "target_agent_type": "agent_a",
+        "status": "COMPLETED",
+        "parent_message_id": "",
+    }
+    exec_b = {
+        "execution_id": "exec-b",
+        "message_id": "msg-b",
+        "session_id": "sess-1",
+        "worker_id": "worker-2",
+        "target_agent_type": "agent_b",
+        "status": "RUNNING",
+        "parent_message_id": "msg-a",
+    }
+
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = [exec_a, exec_b]
+    redis = AsyncMock()
+    client = GatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-1", reason="user abort")
+
+    assert result.success is True
+    assert result.status == "CANCEL_REQUESTED"
+    assert result.cancelled_count == 1
+    assert result.already_finished_count == 1
+
+    # Only the non-terminal execution gets mark_execution_cancelling
+    registry.mark_execution_cancelling.assert_called_once_with(
+        "exec-b", "sess-1", "user abort"
+    )
+    # The terminal execution only gets flagged, not re-marked cancelling
+    registry.mark_cancel_requested.assert_called_once_with(
+        "exec-a", "sess-1", "user abort"
+    )
+
+    # Only one CancelTaskCommand sent, to the non-terminal execution's worker
+    assert redis.xadd.call_count == 1
+    target_streams = [call.args[0] for call in redis.xadd.call_args_list]
+    assert any("worker-2" in s for s in target_streams)
+    assert not any("worker-1" in s for s in target_streams)
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_all_terminal_returns_already_finished():
+    """If every execution in the session is already terminal, cancel_session
+    should report failure with ALREADY_FINISHED and not dispatch anything."""
+    exec_a = {
+        "execution_id": "exec-a",
+        "message_id": "msg-a",
+        "session_id": "sess-1",
+        "worker_id": "worker-1",
+        "target_agent_type": "agent_a",
+        "status": "COMPLETED",
+        "parent_message_id": "",
+    }
+    exec_b = {
+        "execution_id": "exec-b",
+        "message_id": "msg-b",
+        "session_id": "sess-1",
+        "worker_id": "worker-2",
+        "target_agent_type": "agent_b",
+        "status": "FAILED",
+        "parent_message_id": "msg-a",
+    }
+
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = [exec_a, exec_b]
+    redis = AsyncMock()
+    client = GatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-1", reason="user abort")
+
+    assert result.success is False
+    assert result.status == "ALREADY_FINISHED"
+    assert result.cancelled_count == 0
+    assert result.already_finished_count == 2
+
+    registry.mark_execution_cancelling.assert_not_called()
+    assert registry.mark_cancel_requested.call_count == 2
+    redis.xadd.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_empty_session_returns_not_found():
+    """A session with no executions at all should return NOT_FOUND, not be
+    confused with a session whose executions are all already finished."""
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = []
+    redis = AsyncMock()
+    client = GatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-unknown", reason="")
+
+    assert result.success is False
+    assert result.status == "NOT_FOUND"
+    assert result.cancelled_count == 0
+    assert result.already_finished_count == 0
+
+    registry.mark_execution_cancelling.assert_not_called()
+    registry.mark_cancel_requested.assert_not_called()
+    redis.xadd.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_marks_registry_when_worker_not_assigned():
+    """A non-terminal execution that hasn't been claimed by a worker yet
+    (no worker_id) should still be marked cancelling, but no
+    CancelTaskCommand can be routed anywhere for it yet."""
+    exec_a = {
+        "execution_id": "exec-a",
+        "message_id": "msg-a",
+        "session_id": "sess-1",
+        "worker_id": "",
+        "target_agent_type": "agent_a",
+        "status": "QUEUED",
+        "parent_message_id": "",
+    }
+
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = [exec_a]
+    redis = AsyncMock()
+    client = GatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-1", reason="user abort")
+
+    assert result.success is True
+    assert result.status == "CANCEL_REQUESTED"
+    assert result.cancelled_count == 1
+
+    registry.mark_execution_cancelling.assert_called_once_with(
+        "exec-a", "sess-1", "user abort"
+    )
+    redis.xadd.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_client_cancel_session_without_registry_raises():
+    """cancel_session should fail the same way cancel_task does when the
+    client was constructed without a WorkerRegistry."""
+    redis = AsyncMock()
+    client = GatewayClient(registry=None, redis_client=redis)
+
+    with pytest.raises(ValueError):
+        await client.cancel_session(session_id="sess-1")
+
+
+@pytest.mark.asyncio
+async def test_byai_client_cancel_session_inherited_unmodified():
+    """ByaiGatewayClient only overrides send_message; cancel_session should
+    be inherited from GatewayClient and behave identically."""
+    exec_a = {
+        "execution_id": "exec-a",
+        "message_id": "msg-a",
+        "session_id": "sess-1",
+        "worker_id": "worker-1",
+        "target_agent_type": "agent_a",
+        "status": "RUNNING",
+        "parent_message_id": "",
+    }
+
+    registry = AsyncMock()
+    registry.get_all_session_executions.return_value = [exec_a]
+    redis = AsyncMock()
+    client = ByaiGatewayClient(registry=registry, redis_client=redis)
+
+    result = await client.cancel_session(session_id="sess-1", reason="user abort")
+
+    assert result.success is True
+    assert result.status == "CANCEL_REQUESTED"
+    assert result.cancelled_count == 1
+    registry.mark_execution_cancelling.assert_called_once_with(
+        "exec-a", "sess-1", "user abort"
+    )
+    assert redis.xadd.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_client_send_message_resume_reuses_existing_execution_id():
+    """Resuming a suspended execution must reuse its execution_id and must
+    not re-initialize the registry record for it. initialize_execution()
+    unconditionally overwrites the message_id -> execution_id mapping, which
+    would silently detach the ResumeCommand from the original (still
+    WAITING_USER) execution and orphan it."""
+    mock_redis = AsyncMock()
+    mock_registry = AsyncMock()
+    mock_registry.get_execution_by_message_id.return_value = {
+        "execution_id": "exec-original",
+        "message_id": "msg-1",
+        "session_id": "s1",
+        "status": "WAITING_USER",
+    }
+
+    client = GatewayClient(redis_client=mock_redis, registry=mock_registry)
+    await client.send_message(
+        target_agent_type="langgraph_agent",
+        session_id="s1",
+        content="user reply",
+        action_type="RESUME",
+        message_id="msg-1",
+        target_worker_id="worker-42",
+    )
+
+    mock_registry.get_execution_by_message_id.assert_awaited_once_with(
+        "msg-1", session_id="s1"
+    )
+    mock_registry.initialize_execution.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_client_send_message_resume_without_match_falls_back_to_new_execution():
+    """A RESUME whose message_id doesn't resolve to an existing execution
+    (e.g. a genuinely new resume-shaped message) falls back to today's
+    behavior: mint a fresh execution_id and initialize it in the registry."""
+    mock_redis = AsyncMock()
+    mock_registry = AsyncMock()
+    mock_registry.get_execution_by_message_id.return_value = None
+
+    client = GatewayClient(redis_client=mock_redis, registry=mock_registry)
+    await client.send_message(
+        target_agent_type="langgraph_agent",
+        session_id="s1",
+        content="user reply",
+        action_type="RESUME",
+        message_id="msg-unmatched",
+        target_worker_id="worker-42",
+    )
+
+    mock_registry.get_execution_by_message_id.assert_awaited_once_with(
+        "msg-unmatched", session_id="s1"
+    )
+    mock_registry.initialize_execution.assert_awaited_once()
+    initialized = mock_registry.initialize_execution.await_args.args[0]
+    assert initialized["message_id"] == "msg-unmatched"
+    assert initialized["execution_id"]
+
+
+@pytest.mark.asyncio
+async def test_client_send_message_ask_agent_does_not_query_registry_for_resume_match():
+    """A fresh ASK_AGENT dispatch always gets a new message_id, so there is
+    nothing to reattach to -- it should skip the resume lookup entirely and
+    keep minting+initializing a new execution as before."""
+    mock_redis = AsyncMock()
+    mock_registry = AsyncMock()
+
+    client = GatewayClient(redis_client=mock_redis, registry=mock_registry)
+    await client.send_message(
+        target_agent_type="langgraph_agent",
+        session_id="s1",
+        content="hello",
+        action_type="ASK_AGENT",
+        message_id="msg-new",
+        target_worker_id="worker-42",
+    )
+
+    mock_registry.get_execution_by_message_id.assert_not_called()
+    mock_registry.initialize_execution.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_client_send_message_resume_falls_back_when_registry_lacks_lookup():
+    """A registry double/implementation without get_execution_by_message_id
+    must not blow up a RESUME send -- mirror the defensive hasattr() guard
+    already used for initialize_execution() and fall back to minting a new
+    execution_id."""
+    mock_redis = AsyncMock()
+
+    class RegistryWithoutLookup:
+        async def initialize_execution(self, execution):
+            return None
+
+        async def is_worker_online(self, worker_id):
+            return True
+
+    client = GatewayClient(redis_client=mock_redis, registry=RegistryWithoutLookup())
+
+    response = await client.send_message(
+        target_agent_type="langgraph_agent",
+        session_id="s1",
+        content="user reply",
+        action_type="RESUME",
+        message_id="msg-1",
+        target_worker_id="worker-42",
+    )
+
+    assert response.success is True