digitalkin 1.0.5.dev3__tar.gz → 1.0.5.dev4__tar.gz

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 1.0.5.dev3
+Version: 1.0.5.dev4
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/pyproject.toml

@@ -41,7 +41,7 @@
         "pydantic>=2.12.4",
         "redis[hiredis]>=7.4.0,<9",
     ]
-    version = "1.0.5.dev3"
+    version = "1.0.5.dev4"
 
     [project.optional-dependencies]
         agno = [ "agno>=2.6" ]

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "1.0.5.dev3"
+    __version__ = "1.0.5.dev4"

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/community/agno/agno_adapter.py

@@ -676,6 +676,26 @@ class AgnoStreamAdapter:
             self._closed_tool_call_ids.add(tool_call_id)
 
         content = agno_event.content
+        if tool is not None and tool.tool_call_error:
+            # Agno flags a raised tool on the completed event, then sends a tool_call_error
+            # that the closed-id check drops — so the failure is surfaced here.
+            # TODO(validate): AGUI-TOOL-FAILED-STATUS a raised tool surfaces as ToolCallError, not success
+            logger.info(
+                "[VALIDATE AGUI-TOOL-FAILED-STATUS] tool %s (id=%s) failed: %s",
+                tool.tool_name,
+                tool_call_id,
+                tool.result,
+            )
+            return [
+                ToolCallErrorEvent(
+                    event=AgentRunEvent.TOOL_CALL_ERROR,
+                    tool=tool_info,
+                    error_message=str(tool.result) if tool.result else None,
+                    subagent_run_id=self._subagent_of(self._run_key),
+                    timestamp=timestamp,
+                    metadata=self._last_metadata,
+                )
+            ]
         return [
             ToolCallCompletedEvent(
                 event=AgentRunEvent.TOOL_CALL_COMPLETED,

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/community/agno/hitl.py

@@ -11,6 +11,7 @@ import json
 import logging
 from typing import TYPE_CHECKING, Any, ClassVar
 
+from ag_ui.core.events import Interrupt, RunFinishedInterruptOutcome
 from pydantic import BaseModel, ConfigDict, Field
 
 from digitalkin.community.agno.models import PauseInfo
@@ -268,13 +269,14 @@ class HitlEvents:
         run_id: str,
         pending_tool_call_ids: list[str],
     ) -> None:
-        """Emit an AG-UI ``RunFinished`` with ``status="awaiting_tool_result"``.
+        """Emit an AG-UI ``RunFinished`` with an interrupt outcome.
 
         This is the protocol signal telling the front "the run paused on a
-        client-side tool; execute it and reply with a ``ToolMessage``". It
-        goes out via ``context.callbacks.send_message`` (bypassing the
-        standard :class:`~digitalkin.mixins.agui_mixin.AgUiMixin` event
-        mapping, which has no notion of an "awaiting" status).
+        client-side tool; execute it and reply with a ``ToolMessage``". The
+        standard ``outcome={type: "interrupt"}`` carries one interrupt per pending
+        tool; the legacy ``result.status="awaiting_tool_result"`` is kept for
+        existing clients. When an AG-UI run is open, its state first moves the
+        tools to ``awaiting_input`` and the run to ``interrupted``.
 
         Args:
             context: Current module context.
@@ -291,11 +293,31 @@ class HitlEvents:
             AgUiRunFinishedOutput,
         )
 
+        run = context.agui_run
+        tools = run.state["tools"] if run is not None else {}
+        interrupts = [
+            Interrupt(
+                id=tool_call_id,
+                reason="tool_call",
+                tool_call_id=tool_call_id,
+                message=f"Waiting for the result of {tools.get(tool_call_id, {}).get('name') or 'a frontend tool'}",
+            )
+            for tool_call_id in pending_tool_call_ids
+        ]
+        if run is not None and run.open:
+            run.open = False
+            await run.update(
+                context,
+                *((("tools", tid, "status"), "awaiting_input") for tid in pending_tool_call_ids if tid in tools),
+                (("interrupts",), [i.model_dump(mode="json", by_alias=True, exclude_none=True) for i in interrupts]),
+                (("run", "status"), "interrupted"),
+            )
         output = AgUiOutput(
             root=AgUiRunFinishedOutput(
                 event=AgUiRunFinishedEvent(
                     thread_id=thread_id,
                     run_id=run_id,
+                    outcome=RunFinishedInterruptOutcome(interrupts=interrupts) if interrupts else None,
                     result={
                         "status": _AWAITING_STATUS,
                         "pending_tool_call_ids": pending_tool_call_ids,

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/community/agno/module_toolkit.py

@@ -25,6 +25,7 @@ from digitalkin.logger import logger
 from digitalkin.models.module import ModuleContext
 from digitalkin.models.module.ag_ui import AgUiCustomEventOutput, AgUiOutput
 from digitalkin.models.module.tool_cache import ToolDefinition, ToolModuleInfo
+from digitalkin.services.communication.exceptions import ToolCallError
 
 # Default timeout for tool calls in seconds
 DEFAULT_TOOL_TIMEOUT_SECONDS = 300
@@ -265,7 +266,14 @@ class ModuleToolkit(DkToolkit):
         # callbacks is a dict-driven SimpleNamespace (module_context.py:168);
         # send_message may legitimately be absent outside a running job.
         send_message = vars(context.callbacks).get("send_message")
+        value = event.get("value")
+        url = value.get("url") if isinstance(value, dict) else None
         if send_message is None:
+            logger.warning(
+                "Dropped custom event '%s' url=%s: no send_message on the agent context",
+                event["name"],
+                url,
+            )
             return
 
         try:
@@ -275,10 +283,15 @@ class ModuleToolkit(DkToolkit):
             await send_message(
                 AgUiOutput(
                     root=AgUiCustomEventOutput(
-                        event=AgUiCustomEvent(name=event["name"], value=event.get("value")),
+                        event=AgUiCustomEvent(name=event["name"], value=value),
                     )
                 )
             )
+            logger.info(
+                "Relayed custom event '%s' url=%s to the agent stream",
+                event["name"],
+                url,
+            )
         except Exception:
             logger.exception("Failed to relay custom event '%s' to the agent stream", event["name"])
 
@@ -467,6 +480,10 @@ class ModuleToolkit(DkToolkit):
         Wraps the SDK module's async generator function into an async function
         that Agno can consume, with proper error handling and response formatting.
 
+        A failure raises ``ToolCallError`` carrying the same JSON body a success-path caller
+        would have read: Agno hands ``str(exc)`` to the model unchanged and flags the call
+        ``tool_call_error``, which is what marks the tool failed for the front.
+
         Returns:
             Async function that calls the SDK tool and returns a JSON result, or a
             ToolResult when the tool returned images alongside its text output.
@@ -509,14 +526,14 @@ class ModuleToolkit(DkToolkit):
                 duration_ms = round((time.perf_counter() - start_time) * 1000, 2)
                 error_msg = f"Tool '{tool_name}' timed out after {timeout}s"
                 logger.warning("%s task_id=%s", error_msg, task_id)
-                return handle_failure(tool_name, error_msg, duration_ms)
+                raise ToolCallError(handle_failure(tool_name, error_msg, duration_ms))
 
             except Exception as e:
                 outcome = "error"
                 duration_ms = round((time.perf_counter() - start_time) * 1000, 2)
                 error_msg = f"Failed to call tool '{tool_name}': {e!s}"
                 logger.warning("%s task_id=%s", error_msg, task_id, exc_info=True)
-                return handle_failure(tool_name, error_msg, duration_ms)
+                raise ToolCallError(handle_failure(tool_name, error_msg, duration_ms))
 
             else:
                 call_timer.mark("gen_consume")
@@ -536,7 +553,7 @@ class ModuleToolkit(DkToolkit):
 
                 outcome = "no_success"
                 error_msg = ModuleToolkit._extract_error_message(error_frame)
-                return handle_failure(tool_name, error_msg, duration_ms)
+                raise ToolCallError(handle_failure(tool_name, error_msg, duration_ms))
 
             finally:
                 call_timer.mark("respond")

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/core/task_manager/module_runner.py

@@ -21,9 +21,10 @@ from digitalkin.models.core.task_monitor import CancellationReason
 from digitalkin.models.grpc_servers.stream_error_codes import StreamErrorCode
 from digitalkin.models.module.base_types import DataModel
 from digitalkin.models.module.module import ModuleStatus
-from digitalkin.models.module.utility import EndOfStreamOutput, StreamCancelledOutput
+from digitalkin.models.module.utility import EndOfStreamOutput, StreamCancelledOutput, StreamHeartbeatOutput
 from digitalkin.models.services.storage import BaseRole
 from digitalkin.models.settings.gateway import get_gateway_settings
+from digitalkin.models.settings.module import get_module_settings
 from digitalkin.models.settings.profiling import ProfilerMode, get_profiling_settings
 from digitalkin.models.settings.redis import get_redis_settings
 
@@ -143,9 +144,11 @@ class ModuleRunner:
             idem_key = f"idem:{task_id}"
             idem_ttl = get_redis_settings().idem_ttl
             eos_written = False
+            last_write_ns = time.monotonic_ns()
 
             async def _on_output(output_data: Any) -> None:
-                nonlocal first_logged, seq, eos_written
+                nonlocal first_logged, seq, eos_written, last_write_ns
+                last_write_ns = time.monotonic_ns()
                 data = output_data.model_dump(mode="json")
                 if data.get("root", {}).get("protocol") == "stream.end":
                     t_eos_write_start = time.perf_counter_ns()
@@ -259,6 +262,41 @@ class ModuleRunner:
             timer.mark("create_job")
             timer.log("ModuleRunner", task_id)
 
+            async def _heartbeat() -> None:
+                interval_s = get_module_settings().heartbeat_interval_s
+                while not eos_written:
+                    # Sleep only the rest of the silence window, so no gap exceeds one interval.
+                    remaining_s = interval_s - (time.monotonic_ns() - last_write_ns) / 1e9
+                    if remaining_s > 0:
+                        await asyncio.sleep(remaining_s)
+                        continue
+                    # TODO(validate): TASK-HEARTBEAT a silent running task keeps its stream open
+                    logger.info("[VALIDATE TASK-HEARTBEAT] task silent for %.0fs; beating", interval_s, extra=log_extra)
+                    run = module.context.agui_run
+                    try:
+                        if run is not None and run.open:
+                            await run.update(module.context, (("run", "heartbeatAt"), run.now_ms()))
+                        else:
+                            await _on_output(
+                                DataModel[StreamHeartbeatOutput](
+                                    root=StreamHeartbeatOutput(), annotations={"role": BaseRole.SYSTEM}
+                                )
+                            )
+                    except Exception:
+                        # The beat task is never awaited: log here or the failure is lost.
+                        logger.exception("ModuleRunner: heartbeat stopped", extra=log_extra)
+                        return
+
+            module_task = self._servicer.job_manager.tasks.get(job_id)
+            if module_task is None or eos_written:
+                logger.debug(
+                    "ModuleRunner: no heartbeat (task registered=%s)", module_task is not None, extra=log_extra
+                )
+            else:
+                # The task's done-callback holds the only reference and stops the beat with it.
+                beat = asyncio.create_task(_heartbeat())
+                module_task.add_done_callback(lambda _: beat.cancel())
+
             if await self._redis_client.exists(cancel_key):
                 task = self._servicer.job_manager.tasks.get(job_id)
                 session = self._servicer.job_manager.tasks_sessions.get(job_id)

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/mixins/agui_mixin.py

@@ -3,14 +3,16 @@
 This mixin provides utilities to convert framework-agnostic agent events
 into AG-UI protocol events and send them through the module context callbacks.
 
-The mixin is a stateless emitter: it receives events with all necessary info
-(including IDs) and emits the corresponding AG-UI protocol events.
-All state management (ID generation, lifecycle tracking) belongs in the adapter layer.
+The mixin receives events with all necessary info (including IDs) and emits the
+corresponding AG-UI protocol events. Event-sequence bookkeeping belongs in the adapter;
+the only state kept here is the SDK-owned pipeline view (``AgUiRunState``) mirrored to
+the front as ``STATE_SNAPSHOT`` + ``STATE_DELTA``. See ``docs/agui_events.md``.
 """
 
 from __future__ import annotations
 
 import json
+import time
 import uuid
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -23,6 +25,7 @@ from ag_ui.core.events import ReasoningStartEvent as AgUiReasoningStartEvent
 from ag_ui.core.events import RunErrorEvent as AgUiRunErrorEvent
 from ag_ui.core.events import RunFinishedEvent as AgUiRunFinishedEvent
 from ag_ui.core.events import RunStartedEvent as AgUiRunStartedEvent
+from ag_ui.core.events import StateDeltaEvent, StateSnapshotEvent
 from ag_ui.core.events import SubagentErrorEvent as AgUiSubagentErrorEvent
 from ag_ui.core.events import SubagentFinishedEvent as AgUiSubagentFinishedEvent
 from ag_ui.core.events import SubagentStartedEvent as AgUiSubagentStartedEvent
@@ -35,6 +38,7 @@ from ag_ui.core.events import ToolCallResultEvent as AgUiToolCallResultEvent
 from ag_ui.core.events import ToolCallStartEvent as AgUiToolCallStartEvent
 from pydantic import BaseModel
 
+from digitalkin.logger import logger
 from digitalkin.models.events import (
     AgentRunEvent,
     BaseAgentRunEvent,
@@ -67,6 +71,8 @@ from digitalkin.models.module.ag_ui import (
     AgUiRunErrorOutput,
     AgUiRunFinishedOutput,
     AgUiRunStartedOutput,
+    AgUiStateDeltaOutput,
+    AgUiStateSnapshotOutput,
     AgUiSubagentErrorOutput,
     AgUiSubagentFinishedOutput,
     AgUiSubagentStartedOutput,
@@ -84,19 +90,110 @@ if TYPE_CHECKING:
     from digitalkin.models.module.module_context import ModuleContext
 
 
+class AgUiRunState:
+    """SDK-owned pipeline state of one task's AG-UI run.
+
+    The front only renders it: ``snapshot`` replaces its state, ``update`` sends the
+    matching JSON Patch. Shape: ``run`` (status, error, heartbeatAt), ``tools`` and
+    ``subagents`` keyed by id (name, status, startedAt, error), ``interrupts``.
+    """
+
+    def __init__(self) -> None:
+        """Start a running run with an empty pipeline."""
+        self.open = True
+        self.state: dict[str, Any] = {
+            "run": {"status": "running", "error": None, "heartbeatAt": None},
+            "tools": {},
+            "subagents": {},
+            "interrupts": [],
+        }
+
+    async def snapshot(self, context: ModuleContext) -> None:
+        """Emit the full state as ``STATE_SNAPSHOT``.
+
+        Args:
+            context: Module context whose callbacks carry the event.
+        """
+        await context.callbacks.send_message(
+            AgUiOutput(root=AgUiStateSnapshotOutput(event=StateSnapshotEvent(snapshot=self.state)))
+        )
+
+    async def update(self, context: ModuleContext, *changes: tuple[tuple[str, ...], Any]) -> None:
+        """Set each ``(key path, value)`` in the state and emit them as one ``STATE_DELTA``.
+
+        Args:
+            context: Module context whose callbacks carry the event.
+            *changes: Key path into the state and the value to set there.
+        """
+        ops = []
+        for keys, value in changes:
+            node = self.state
+            for key in keys[:-1]:
+                node = node[key]
+            ops.append({
+                "op": "replace" if keys[-1] in node else "add",
+                "path": "/" + "/".join(key.replace("~", "~0").replace("/", "~1") for key in keys),
+                "value": value,
+            })
+            node[keys[-1]] = value.copy() if isinstance(value, dict | list) else value
+        await context.callbacks.send_message(AgUiOutput(root=AgUiStateDeltaOutput(event=StateDeltaEvent(delta=ops))))
+
+    async def fail(
+        self,
+        context: ModuleContext,
+        status: str,
+        message: str,
+        code: str | None,
+        **fields: Any,
+    ) -> None:
+        """Close the run once: record ``status``/``message`` in state, then emit ``RUN_ERROR``.
+
+        Args:
+            context: Module context whose callbacks carry the events.
+            status: Terminal run status, ``failed`` or ``cancelled``.
+            message: Human-readable error.
+            code: ``RUN_ERROR`` code.
+            **fields: Extra AG-UI event fields (metadata).
+        """
+        if not self.open:
+            logger.info(
+                "AG-UI run already closed (%s); dropping %s RUN_ERROR code=%s message=%s",
+                self.state["run"]["status"],
+                status,
+                code,
+                message,
+                extra=context.session.current_ids(),
+            )
+            return
+        self.open = False
+        await self.update(context, (("run", "status"), status), (("run", "error"), message))
+        await context.callbacks.send_message(
+            AgUiOutput(root=AgUiRunErrorOutput(event=AgUiRunErrorEvent(message=message, code=code, **fields)))
+        )
+
+    @staticmethod
+    def now_ms() -> int:
+        """Wall-clock epoch milliseconds, the unit of every ``*At`` field.
+
+        Returns:
+            Current time in epoch milliseconds.
+        """
+        return time.time_ns() // 1_000_000
+
+
 class AgUiMixin:
     """Mixin for converting agent events to AG-UI protocol and sending them.
 
-    This mixin is a stateless emitter: each handler reads IDs from the event
-    and emits the corresponding AG-UI event(s). The adapter is responsible for
-    generating IDs and managing event lifecycle (start/complete sequences).
+    Each handler reads IDs from the event and emits the corresponding AG-UI event(s),
+    plus the ``STATE_DELTA`` that keeps ``context.agui_run`` in sync. The adapter is
+    responsible for generating IDs and managing event lifecycle (start/complete sequences).
 
     Usage::
 
         class MyTrigger(BaseTrigger, AgUiMixin):
             async def execute(self, context, input_data):
                 async for event in agent.run(input_data.message, stream=True):
-                    await self.agui_send_message(context, event)
+                    await self.send_message(context, event)
     """
 
     def __init__(self) -> None:
@@ -217,6 +314,8 @@ class AgUiMixin:
             )
         )
         await self._send_agui(context, output)
+        context.agui_run = AgUiRunState()
+        await context.agui_run.snapshot(context)
 
     async def _handle_text_message_started(
         self,
@@ -285,6 +384,10 @@ class AgUiMixin:
             event.metadata,
             extra=context.session.current_ids(),
         )
+        run = context.agui_run
+        if run is not None and run.open:
+            run.open = False
+            await run.update(context, (("run", "status"), "completed"))
         output = AgUiRunFinishedOutput(
             event=AgUiRunFinishedEvent(
                 thread_id=self._thread_id,
@@ -299,8 +402,11 @@ class AgUiMixin:
         context: ModuleContext,
         event: RunErrorEvent,
     ) -> None:
-        """Handle run error event - emit AG-UI RunError."""
+        """Handle run error event - mark the run failed and emit AG-UI RunError."""
         error_msg = event.content or "Agent run failed"
+        if context.agui_run is not None:
+            await context.agui_run.fail(context, "failed", error_msg, event.error_type, **self._authored(event))
+            return
         output = AgUiRunErrorOutput(
             event=AgUiRunErrorEvent(
                 message=error_msg,
@@ -326,6 +432,14 @@ class AgUiMixin:
             )
         )
         await self._send_agui(context, output)
+        if context.agui_run is not None and event.subagent_run_id:
+            await context.agui_run.update(
+                context,
+                (
+                    ("subagents", event.subagent_run_id),
+                    {"name": event.name, "status": "running", "startedAt": AgUiRunState.now_ms(), "error": None},
+                ),
+            )
 
     async def _handle_subagent_finished(
         self,
@@ -341,6 +455,9 @@ class AgUiMixin:
             )
         )
         await self._send_agui(context, output)
+        run = context.agui_run
+        if run is not None and event.subagent_run_id is not None and event.subagent_run_id in run.state["subagents"]:
+            await run.update(context, (("subagents", event.subagent_run_id, "status"), "completed"))
 
     async def _handle_subagent_error(
         self,
@@ -361,6 +478,13 @@ class AgUiMixin:
             )
         )
         await self._send_agui(context, output)
+        run = context.agui_run
+        if run is not None and event.subagent_run_id is not None and event.subagent_run_id in run.state["subagents"]:
+            await run.update(
+                context,
+                (("subagents", event.subagent_run_id, "status"), "failed"),
+                (("subagents", event.subagent_run_id, "error"), event.message),
+            )
 
     async def _handle_tool_call_started(
         self,
@@ -394,6 +518,21 @@ class AgUiMixin:
             )
             await self._send_agui(context, args_output)
 
+        if context.agui_run is not None:
+            await context.agui_run.update(
+                context,
+                (
+                    ("tools", tool_call_id),
+                    {
+                        "name": tool.tool_name,
+                        "status": "running",
+                        "subagentRunId": event.subagent_run_id,
+                        "startedAt": AgUiRunState.now_ms(),
+                        "error": None,
+                    },
+                ),
+            )
+
     async def _handle_tool_call_completed(
         self,
         context: ModuleContext,
@@ -425,19 +564,46 @@ class AgUiMixin:
             )
             await self._send_agui(context, result_output)
 
+        run = context.agui_run
+        if run is not None and tool_call_id in run.state["tools"]:
+            await run.update(context, (("tools", tool_call_id, "status"), "completed"))
+
     async def _handle_tool_call_error(
         self,
         context: ModuleContext,
         event: ToolCallErrorEvent,
     ) -> None:
-        """Handle tool call error event - emit AG-UI ToolCallEnd."""
+        """Handle tool call error event - emit AG-UI ToolCallEnd + an error ToolCallResult, mark it failed.
+
+        ``TOOL_CALL_RESULT`` has no error flag, so the failure is carried by its
+        ``{"error": ...}`` content and by ``tools/<id>/status = failed`` in state.
+        """
         tool = event.tool
         if not tool:
             return
 
         tool_call_id = tool.tool_call_id or str(uuid.uuid4())
+        error_msg = event.error_message or "Tool call failed"
         output = AgUiToolCallEndOutput(event=AgUiToolCallEndEvent(tool_call_id=tool_call_id, **self._authored(event)))
         await self._send_agui(context, output)
+        result_output = AgUiToolCallResultOutput(
+            event=AgUiToolCallResultEvent(
+                message_id=str(uuid.uuid4()),
+                tool_call_id=tool_call_id,
+                content=json.dumps({"error": error_msg}),
+                role="tool",
+                **self._authored(event),
+            )
+        )
+        await self._send_agui(context, result_output)
+
+        run = context.agui_run
+        if run is not None and tool_call_id in run.state["tools"]:
+            await run.update(
+                context,
+                (("tools", tool_call_id, "status"), "failed"),
+                (("tools", tool_call_id, "error"), error_msg),
+            )
 
     async def _handle_reasoning_started(
         self,

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/models/module/module_context.py

@@ -4,7 +4,7 @@ import asyncio
 from collections.abc import AsyncGenerator, Callable
 from datetime import tzinfo
 from types import SimpleNamespace
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from zoneinfo import ZoneInfo
 
 from google.protobuf import json_format
@@ -28,6 +28,9 @@ from digitalkin.services.storage.storage_strategy import StorageStrategy
 from digitalkin.services.task_manager.task_manager_strategy import TaskManagerStrategy
 from digitalkin.services.user_profile.user_profile_strategy import UserProfileStrategy
 
+if TYPE_CHECKING:
+    from digitalkin.mixins.agui_mixin import AgUiRunState
+
 
 class Session(SimpleNamespace):
     """Session data container with mandatory setup_id and mission_id."""
@@ -116,6 +119,7 @@ class ModuleContext:
     shared: dict[str, Any]
     tool_cache: ToolCache
     request_metadata: RequestMetadata
+    agui_run: "AgUiRunState | None"
 
     def __init__(  # All service strategies are mandatory constructor args # ruff: ignore[too-many-arguments, too-many-positional-arguments]
         self,
@@ -180,6 +184,7 @@ class ModuleContext:
         self.shared = shared if shared is not None else {}
         self.tool_cache = tool_cache or ToolCache()
         self.request_metadata = RequestMetadata(request_metadata)
+        self.agui_run = None
 
     async def get_module_schemas_by_id(
         self,

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/models/module/utility.py

@@ -34,6 +34,12 @@ class StreamCancelledOutput(UtilityProtocol):
     reason: str = "cancelled"
 
 
+class StreamHeartbeatOutput(UtilityProtocol):
+    """Liveness beat from a module that is running but silent; consumers ignore it."""
+
+    protocol: Literal["stream.heartbeat"] = "stream.heartbeat"
+
+
 class HealthcheckPingInput(UtilityProtocol):
     """Input for healthcheck ping request."""
 

{digitalkin-1.0.5.dev3 → digitalkin-1.0.5.dev4}/src/digitalkin/models/settings/module.py

@@ -32,6 +32,14 @@ class ModuleSettings(BaseSettings):
     file_history_flush_threshold: int = Field(
         default=10, description="Dirty-entry count that triggers a file-history flush."
     )
+    heartbeat_interval_s: float = Field(
+        default=30.0,
+        gt=0,
+        description=(
+            "Silence after which a running task writes a liveness frame (AG-UI STATE_DELTA or "
+            "stream.heartbeat). Keep it below the gateway read/dial-back idle timeouts."
+        ),
+    )
 
 
 @lru_cache(maxsize=1)