feat/sub_stage_propagation

docs/designs/sub-stage-propagation.md

@@ -0,0 +1,160 @@
+# Sub-Stage Propagation
+
+**Status:** Approved  
+**Approved:** 2026-09-04  
+**Author:** Aleksei Korota
+
+---
+
+## Problem
+
+When a QuickApp calls another QuickApp as a tool, the sub-app's UI stages are silently discarded — the user sees only an empty "Calling X" stage with no indication of what the sub-app is actually doing.
+
+---
+
+## Root Cause
+
+`DialCompletionService._consume_stream()` constructs `ChatStreamConfig` without a `destination` and without `propagate_stages=True`. As a result, `ChoiceUiSink` is instantiated but completely inert — its `on_delta()` returns immediately when `destination is None`. The sub-app's stage deltas are accumulated into `ChatStreamAccumulator.stages` by `AccumulationSink`, but `complete_request_async()` never reads that field, so the stages are effectively thrown away.
+
+```
+BaseDeploymentTool
+  └─ DialCompletionService._consume_stream()
+       └─ ChatStreamConfig(stage_wrapper=...) ← no destination, no propagate_stages
+            └─ ChoiceUiSink(destination=None)  ← inert, sub-app stages discarded
+```
+
+The sub-app's **text content** does reach the parent (via `StageWrapperUiSink`), so the "Calling X" stage renders the response text. Only stage deltas from the sub-app's `custom_content.stages` are lost.
+
+---
+
+## Approach 1 — Flat Propagation (short-term)
+
+Sub-app stages are re-emitted into the parent's flat stage list with a name prefix that indicates their origin. No SDK or UI changes required.
+
+### Mechanism
+
+1. Inject `Choice` (already request-scoped) into `DialCompletionService`.
+2. Pass it as `destination` and set `propagate_stages=True` in `_consume_stream`.
+3. Thread a `sub_stage_prefix` string (e.g. `"[WeatherApp]"`) from `BaseDeploymentTool` through `ChatStreamConfig` to `ChoiceUiSink._stream_stage_delta()`, where it is prepended to the stage name on creation.
+
+For deep nesting (A → B → C), C's stages reach A's stream as siblings of B's stages, all with their respective prefixes. The feature is gated by a new `PreviewField` on `ApplicationConfig` (`orchestrator.propagate_sub_stages`, default `true`) and by the existing `ENABLE_PREVIEW_FEATURES` env switch.
+
+### Stage index safety
+
+`ChoiceUiSink` tracks sub-app stage indices in a local `_stages_by_index` dict. Actual allocation is done through `destination.create_stage()`, which increments `Choice._last_stage_index` independently — no collision with the parent's own stages.
+
+### Files changed
+
+| File | Change |
+|------|--------|
+| `config/application.py` | `PreviewField` `orchestrator.propagate_sub_stages: bool` |
+| `dial_deployment_tooling/dial_completion_service.py` | Inject `Choice`; pass `destination` + prefix to `_consume_stream` |
+| `dial_deployment_tooling/base_deployment_tool.py` | Pass `application_name` as prefix to `complete_request_async` |
+| `common/chat_completion_stream/handler.py` | Add `sub_stage_prefix: str | None` to `ChatStreamConfig` |
+| `common/chat_completion_stream/chat_stream_sink_factory.py` | Forward `sub_stage_prefix` to `ChoiceUiSink` |
+| `common/chat_completion_stream/choice_ui_stream_sink.py` | Prepend prefix in `_stream_stage_delta` |
+
+### Trade-offs
+
+**Pros:** No SDK bump, no UI changes, deployable today.  
+**Cons:** Flat list without structural hierarchy; deep nesting produces a long undifferentiated list.
+
+---
+
+## Approach 2 — Nested Stages (target architecture)
+
+Sub-app stages are re-emitted as children of the open "Calling X" stage, preserving the call hierarchy at any depth. Requires changes to `aidial_sdk` and a wire-format contract that the UI team must implement.
+
+### Wire format change (`aidial_sdk`)
+
+A new optional field `parent_stage_index` is added to `StartStageChunk`. When present, the UI renders the stage as a child of the referenced stage rather than a top-level sibling.
+
+**Before:**
+```json
+{"custom_content": {"stages": [
+  {"index": 0, "name": "Calling WeatherApp", "status": null},
+  {"index": 1, "name": "Fetching forecast",  "status": null}
+]}}
+```
+
+**After:**
+```json
+{"custom_content": {"stages": [
+  {"index": 0, "name": "Calling WeatherApp", "status": null},
+  {"index": 1, "name": "Fetching forecast", "parent_stage_index": 0, "status": null}
+]}}
+```
+
+`parent_stage_index` is omitted for top-level stages (backwards-compatible). `FinishStageChunk`, `ContentStageChunk`, `NameStageChunk`, and `AttachmentStageChunk` are identified by `index` alone and require no changes.
+
+### SDK API
+
+Two additions to `aidial_sdk`:
+
+```python
+# Option A — factory method on Stage
+child = parent_stage.create_child_stage("Fetching forecast")
+
+# Option B — optional parameter on Choice.create_stage
+child = choice.create_stage("Fetching forecast", parent=parent_stage)
+```
+
+`Stage._stage_index` is exposed as a read-only property `Stage.stage_index`. `StartStageChunk.__init__` accepts `parent_stage_index: int | None = None` and includes it in `to_dict()` via `remove_nones`.
+
+### Mechanism (backend)
+
+1. `BaseStageWrapper` adds `@property stage -> Stage` to expose its private `self.__stage`.
+2. `BaseDeploymentTool._run_in_stage_async()` passes `stage_wrapper.stage` to `complete_request_async()` as a new `parent_stage: Stage | None` parameter.
+3. `DialCompletionService._consume_stream()` forwards `parent_stage` in `ChatStreamConfig`.
+4. `ChoiceUiSink` receives `parent_stage` and uses it when re-emitting sub-app stages:
+   - Top-level sub-stages (no `parent_stage_index` in delta): created with `parent=parent_stage`.
+   - Nested sub-stages (`parent_stage_index` present): remapped through the local `_stages_by_index` dict to find the `Stage` object created in the parent context, then passed as `parent=`.
+
+The `_stages_by_index` dict (already maintained by `ChoiceUiSink`) naturally handles arbitrary recursion depth: each level's index space is remapped independently as it is processed.
+
+The feature is gated identically to Approach 1 (`PreviewField` + `ENABLE_PREVIEW_FEATURES`).
+
+### UI contract (for the UI team)
+
+The UI must:
+- Accept `parent_stage_index: number | undefined` on each stage object in `custom_content.stages`.
+- Build a tree from stage deltas as they arrive in the stream: when `parent_stage_index` is set, attach the stage as a child of the already-opened stage at that index.
+- Render nested stages as collapsible sections inside their parent stage.
+- Treat absent `parent_stage_index` as a top-level stage (no behaviour change for existing flat streams).
+
+### Files changed
+
+**`aidial_sdk`:**
+
+| File | Change |
+|------|--------|
+| `chat_completion/chunks.py` | `StartStageChunk`: add `parent_stage_index: int | None`; update `to_dict` |
+| `chat_completion/stage.py` | Add `parent_stage_index` param; expose `stage_index` property |
+| `chat_completion/choice.py` | `create_stage()`: add `parent: Stage | None = None` |
+
+**`quickapps-backend`:**
+
+| File | Change |
+|------|--------|
+| `config/application.py` | `PreviewField` `orchestrator.propagate_sub_stages: bool` |
+| `common/_stage_delta_types.py` | Add `parent_stage_index` to `StageDeltaItem` TypedDict |
+| `common/base_stage_wrapper.py` | Add `@property stage -> Stage` |
+| `dial_deployment_tooling/base_deployment_tool.py` | Pass `stage_wrapper.stage` as `parent_stage` |
+| `dial_deployment_tooling/dial_completion_service.py` | Add `parent_stage` param; forward to `_consume_stream` |
+| `common/chat_completion_stream/handler.py` | Add `parent_stage: Stage | None` to `ChatStreamConfig` |
+| `common/chat_completion_stream/chat_stream_sink_factory.py` | Forward `parent_stage` to `ChoiceUiSink` |
+| `common/chat_completion_stream/choice_ui_stream_sink.py` | Index remapping + `parent=` on `create_stage` |
+
+### Trade-offs
+
+**Pros:** Correct structural representation; scales cleanly to arbitrary nesting depth.  
+**Cons:** Requires SDK version bump and UI-team implementation of tree rendering.
+
+---
+
+## Scale of changes
+
+| | `aidial_sdk` | `quickapps-backend` |
+|---|---|---|
+| Approach 1 | — | 6 files, ~40 lines |
+| Approach 2 | 3 files, ~40 lines | 8 files, ~80 lines |

docs/generated-app-schema.json

@@ -3188,6 +3188,20 @@
           "$ref": "#/$defs/StageDisplayLevel",
           "default": "info",
           "description": "Threshold for stage visibility. none=no stages at all; errors=failures only; info=user-facing (default); debug=all."
+        },
+        "propagate_sub_stages": {
+          "anyOf": [
+            {
+              "type": "boolean"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "description": "When True (default when preview features are enabled), stages emitted by sub-apps called as tools are propagated as nested children of the 'Calling X' stage. Set to False to disable.",
+          "title": "Propagate Sub Stages",
+          "x-preview": true
         }
       },
       "title": "StageDisplayConfig",

src/quickapp/common/_stage_delta_types.py

@@ -17,6 +17,7 @@ class StageDeltaItem(TypedDict, total=False):
     content: str
     attachments: list[dict[str, Any]]
     status: str
+    parent_stage_index: int
 
 
 def as_stage_delta(item: dict[str, Any]) -> StageDeltaItem:

src/quickapp/common/base_stage_wrapper.py

@@ -31,6 +31,10 @@ def __init__(
             extract_parameters_config_map(tool_config)
         )
 
+    @property
+    def stage(self) -> Stage:
+        return self.__stage
+
     def __enter__(self) -> "BaseStageWrapper":
         if self.__already_open:
             # Stage was opened while tool-call arguments streamed; skip re-open.

src/quickapp/common/chat_completion_stream/accumulation_stream_sink.py

@@ -9,9 +9,8 @@
 class AccumulationSink(ChatStreamSink):
     """Always active: builds the in-memory stream result for history / execute / logs."""
 
-    def __init__(self, accumulator: ChatStreamAccumulator, *, stream_content: bool = True) -> None:
+    def __init__(self, accumulator: ChatStreamAccumulator) -> None:
         self._accumulator = accumulator
-        self._stream_content = stream_content
 
     def on_stream_start(self) -> None:
         return
@@ -25,7 +24,7 @@ def on_delta(self, delta: NormalizedChoiceDelta) -> None:
             if norm.state is not None:
                 self._accumulator.merge_state(norm.state)
 
-        if delta.content and self._stream_content:
+        if delta.content:
             self._accumulator.append_content(delta.content)
 
         for tool_call in delta.tool_calls:

src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py

@@ -81,13 +81,14 @@ def create(self, config: ChatStreamConfig) -> ChatStreamPipeline:
         return ChatStreamPipeline(
             accumulator,
             [
-                AccumulationSink(accumulator, stream_content=config.stream_content),
+                AccumulationSink(accumulator),
                 ChoiceUiSink(
                     accumulator,
                     destination=config.destination,
                     stream_content=config.stream_content,
                     propagate_stages=config.propagate_stages,
                     tools_by_name=_tools_by_name(self._tools()),
+                    parent_stage=config.parent_stage,
                 ),
                 StageWrapperUiSink(
                     stage_wrapper=config.stage_wrapper,

src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py

@@ -59,6 +59,7 @@ def __init__(
         stream_content: bool = True,
         propagate_stages: bool = False,
         tools_by_name: dict[str, StagedBaseTool] | None = None,
+        parent_stage: Stage | None = None,
     ) -> None:
         self._accumulator = accumulator
         self._destination = destination
@@ -68,6 +69,7 @@ def __init__(
         self._stages_by_index: dict[int, Stage] = {}
         self._tool_stages_by_index: dict[int, _StreamingToolStageState] = {}
         self._suppressed_tool_indexes: set[int] = set()
+        self._parent_stage: Stage | None = parent_stage
 
     def on_stream_start(self) -> None:
         destination = self._destination
@@ -114,7 +116,7 @@ def on_stream_failure(self) -> None:
     def _apply_custom(self, norm: NormalizedCustomContent) -> None:
         destination = self._destination
         assert destination is not None
-        if norm.attachments:
+        if norm.attachments and self._parent_stage is None:
             self._add_attachments(destination, norm.attachments)
         for position, raw in norm.stage_entries:
             stage_delta = as_stage_delta(raw)
@@ -314,6 +316,23 @@ def _publish_adopted_tool_stages(self) -> None:
             )
             self._tool_stages_by_index.pop(index, None)
 
+    def _resolve_parent(self, delta: StageDeltaItem) -> Stage | None:
+        """Resolve the parent Stage for a new sub-stage, or None for flat behaviour."""
+        if self._parent_stage is None:
+            return None
+        sub_parent_idx = delta.get("parent_stage_index")
+        if sub_parent_idx is not None:
+            parent = self._stages_by_index.get(sub_parent_idx)
+            if parent is None:
+                logger.warning(
+                    "parent_stage_index=%s not yet seen when creating index=%s; "
+                    "falling back to top-level parent",
+                    sub_parent_idx,
+                    delta.get("index"),
+                )
+            return parent if parent is not None else self._parent_stage
+        return self._parent_stage
+
     def _stream_stage_delta(self, delta: StageDeltaItem, position: int) -> None:
         destination = self._destination
         assert destination is not None
@@ -331,8 +350,9 @@ def _stream_stage_delta(self, delta: StageDeltaItem, position: int) -> None:
                 )
                 log_payload(logger, "Stage delta with missing name: %s", delta)
                 return
+            parent = self._resolve_parent(delta)
             try:
-                stage = destination.create_stage(stage_name)
+                stage = destination.create_stage(stage_name, parent=parent)  # type: ignore[call-arg]
                 stage.open()
                 self._stages_by_index[idx] = stage
                 just_created = True

src/quickapp/common/chat_completion_stream/handler.py

@@ -1,7 +1,7 @@
 import logging
 from collections.abc import AsyncIterable
 
-from aidial_sdk.chat_completion import Choice
+from aidial_sdk.chat_completion import Choice, Stage
 from injector import inject
 from openai import APIError, BadRequestError
 from openai.types.chat import ChatCompletionChunk
@@ -36,6 +36,7 @@ class ChatStreamConfig(BaseModel):
     stage_wrapper: BaseStageWrapper | None = None
     stream_content: bool = True
     propagate_stages: bool = False
+    parent_stage: Stage | None = None
 
 
 class ChatCompletionStreamHandler:

src/quickapp/config/application.py

@@ -42,6 +42,14 @@ class StageDisplayConfig(BaseModel):
         default=StageDisplayLevel.INFO,
         description="Threshold for stage visibility. none=no stages at all; errors=failures only; info=user-facing (default); debug=all.",
     )
+    propagate_sub_stages: bool | None = PreviewField(  # type: ignore[assignment]
+        default=None,
+        description=(
+            "When True (default when preview features are enabled), stages emitted by "
+            "sub-apps called as tools are propagated as nested children of the "
+            "'Calling X' stage. Set to False to disable."
+        ),
+    )
 
 
 def get_max_iterations() -> int:

src/quickapp/dial_deployment_tooling/base_deployment_tool.py

@@ -20,6 +20,7 @@
 from quickapp.common.payload_logging import log_payload
 from quickapp.common.perf_timer.perf_timer import PerformanceTimer
 from quickapp.common.utils import to_plain_dict
+from quickapp.config.application import ApplicationConfig
 from quickapp.config.dial_deployment import DialDeploymentParameters
 from quickapp.config.tools.base import ConfigurableSchemaSimpleType, JsonTypeEnum, OpenAiToolConfig
 from quickapp.config.tools.deployment import ContentPropagation, DialDeploymentTool
@@ -50,6 +51,7 @@ def __init__(
         messages_mixin: MessagesMixin,
         perf_timer: PerformanceTimer,
         stage_wrapper_builder: AssistedBuilder[DeploymentStageWrapper],
+        app_config: ApplicationConfig,
         argument_transformers: list[ToolArgumentTransformer] | None = None,
         **kwargs: Any,
     ):
@@ -65,6 +67,7 @@ def __init__(
         self.__dial_completion_service: DialCompletionService = dial_completion_service
         self.__attachment_resolver: AttachmentResolver = attachment_resolver
         self.__content_propagation: ContentPropagation | None = content_propagation
+        self.__app_config: ApplicationConfig = app_config
         if content_propagation and content_propagation.propagate_history:
             logger.warning(
                 "The 'propagate_history' parameter is deprecated and will be removed in a future release. "
@@ -123,6 +126,11 @@ async def _run_in_stage_async(
         tool_config = cast(DialDeploymentTool, self.tool_config)
         session_id, is_first_call = self._setup_session(kwargs, tool_config, tool_call_id)
         history = await self._resolve_history(tool_config, session_id)
+        stage_display = (
+            self.__app_config.features.stage_display if self.__app_config.features else None
+        )
+        propagate = stage_display.propagate_sub_stages is not False if stage_display else True
+        parent_stage = stage_wrapper.stage if (propagate and stage_wrapper is not None) else None
         result = await self.__dial_completion_service.complete_request_async(
             kwargs,
             self.__application_id,
@@ -131,6 +139,7 @@ async def _run_in_stage_async(
             attachment_urls,
             history=history,
             supports_url_attachments=tool_config.supports_url_attachments,
+            parent_stage=parent_stage,
         )
         if is_first_call and session_id:
             result.content = result.content + f"\n\n[session_id: {session_id}]"

src/quickapp/dial_deployment_tooling/deployment_tool.py

@@ -4,7 +4,7 @@
 from quickapp.common.abstract.base_tool_argument_transformer import ToolArgumentTransformer
 from quickapp.common.messages_mixin import MessagesMixin
 from quickapp.common.perf_timer.perf_timer import PerformanceTimer
-from quickapp.config.application import StageDisplayLevel
+from quickapp.config.application import ApplicationConfig, StageDisplayLevel
 from quickapp.config.tools.deployment import ContentPropagation, DialDeploymentTool
 
 from ._attachment_resolver import AttachmentResolver
@@ -32,6 +32,7 @@ def __init__(
         messages_mixin: MessagesMixin,
         stage_wrapper_builder: AssistedBuilder[DeploymentStageWrapper],
         perf_timer: PerformanceTimer,
+        app_config: ApplicationConfig,
         stage_display_level: StageDisplayLevel = StageDisplayLevel.INFO,
         argument_transformers: list[ToolArgumentTransformer] | None = None,
     ):
@@ -46,6 +47,7 @@ def __init__(
             stage_wrapper_builder=stage_wrapper_builder,
             description=description,
             perf_timer=perf_timer,
+            app_config=app_config,
             stage_display_level=stage_display_level,
             argument_transformers=argument_transformers,
         )