Prompt file imported from msft-ariel-meistel/automl-agentic-platform (.github/prompts/track-b-dts-parity.prompt.md). Copyright stays with the author.

Track B — DTS Feature Parity

The Durable Task Scheduler (DTS) orchestrator (runtime/runtime/dts/orchestrator.py) is the production execution path — it runs when DTS_ENDPOINT + DTS_TASKHUB are configured (the Container App always has them set). However, the DTS path currently lacks five capabilities that the in-process orchestrator (runtime/runtime/orchestrator/__init__.py) already implements.

When DTS is enabled, the API dispatches via dts_client.schedule_run() which only forwards (run_id, tenant_id, goal, nodes). The in-process fallback path passes memory_assembler, working_memory, thread_mirror, and sb_publisher — all of which are silently skipped on the DTS path.

Read AGENTS.md and .github/instructions/python.instructions.md before writing any code.

Files to read first

• runtime/runtime/api.py — lines 718-746 (DTS vs in-process dispatch)
• runtime/runtime/orchestrator/__init__.py — the in-process orchestrator (reference implementation)
• runtime/runtime/dts/orchestrator.py — the DTS orchestrator (target for parity)
• runtime/runtime/dts/client.py — DTS scheduling client
• runtime/runtime/dts/worker.py — DTS worker lifecycle
• runtime/runtime/dts/projection.py — Cosmos projection writer
• runtime/runtime/orchestrator/servicebus_publisher.py — SB publisher Protocol
• runtime/runtime/orchestrator/thread_mirror.py — thread mirror Protocol
• runtime/runtime/orchestrator/working_memory_writer.py — working memory Protocol (if exists)

Gap 1 — Merge-node input semantics (P1)

In-process behavior: _build_user_input() in orchestrator/__init__.py builds labelled merge blocks for nodes with multiple dependencies:

[Upstream outputs follow. Treat them as context for your task.]

From node-1:
<output from node-1>

From node-2:
<output from node-2>

[Your task:]

<task_instructions>

It also chains previous_response_id to the earliest root ancestor.

DTS gap: dag_orchestrator() at line ~294 passes only task_instructions as-is in NodeInput. Merge nodes receive raw task instructions without upstream outputs, breaking multi-agent workflows.

Fix:

1. Add a _build_dts_user_input() function in dts/orchestrator.py that mirrors _build_user_input() logic but uses node_results (the DTS result dict) instead of RunState.node_states:

   def _build_dts_user_input(
       node: dict[str, Any],
       nodes_by_id: dict[str, dict[str, Any]],
       node_results: dict[str, dict[str, Any]],
   ) -> tuple[str, str | None]:
       """Build user input for a node, merging upstream outputs for
       nodes with multiple dependencies. Mirrors the in-process
       orchestrator's _build_user_input() semantics."""
       deps = node.get("depends_on", [])
       task_instructions = node.get("task_instructions", "")
   
       if not deps:
           return task_instructions, None
       if len(deps) == 1:
           parent_result = node_results.get(deps[0], {})
           return task_instructions, parent_result.get("response_id")
   
       # Merge: prepend each parent's output
       blocks = [
           "[Upstream outputs follow. Treat them as context for your task.]",
           "",
       ]
       for dep_id in deps:
           parent_result = node_results.get(dep_id, {})
           blocks.append(f"From {dep_id}:")
           blocks.append(parent_result.get("output_text", "") or "")
           blocks.append("")
       blocks.append("[Your task:]")
       blocks.append("")
       blocks.append(task_instructions)
   
       # Chain to earliest root ancestor
       ancestor_response = None
       cur = node
       seen: set[str] = set()
       while cur.get("depends_on"):
           if cur["id"] in seen:
               break
           seen.add(cur["id"])
           parent = nodes_by_id.get(cur["depends_on"][0])
           if parent is None:
               break
           cur = parent
       root_result = node_results.get(cur["id"], {})
       ancestor_response = root_result.get("response_id")
   
       return "\n".join(blocks), ancestor_response

2. In dag_orchestrator(), replace the current NodeInput construction with:

   user_input, prev_response_id = _build_dts_user_input(
       node, nodes_by_id, node_results,
   )
   
   node_input = NodeInput(
       ...
       task_instructions=user_input,
       previous_response_id=prev_response_id,
       ...
   )

3. Tests:

   • Add test_dts_merge_node_receives_upstream_outputs — two predecessor nodes complete, merge node should receive labelled upstream outputs in its task_instructions.
   • Add test_dts_single_dep_chains_response_id — one dependency, verify previous_response_id is set.

Gap 2 — Memory assembler in DTS activities (P2)

In-process behavior: _run_node() calls memory_assembler.assemble() to prepend a memory block to the agent input. This is best-effort — failures are swallowed.

DTS gap: execute_node_activity() has no memory assembly.

Fix:

1. Add a module-level _memory_assembler reference (same pattern as _invoker and _cosmos_writer):

   _memory_assembler: Any | None = None
   
   def configure_memory_assembler(assembler: Any) -> None:
       global _memory_assembler
       _memory_assembler = assembler

2. In execute_node_activity(), before invoking the agent, prepend memory context (best-effort):

   memory_block = ""
   if _memory_assembler is not None:
       try:
           ctx = await _memory_assembler.assemble(
               tenant_id=node_input.tenant_id,
               user_id="poc-user",
               agent_id=node_input.agent_name,
               run_id=node_input.run_id,
           )
           if ctx is not None:
               memory_block = (
                   "[Memory context — facts from prior runs:]\n"
                   + ctx.text
                   + "\n[End memory context]\n\n"
               )
       except Exception:
           pass
   call.user_input = memory_block + call.user_input

3. In api.py lifespan, call configure_memory_assembler() alongside configure_invoker() and configure_cosmos_writer().

4. Tests:

   • Add test_dts_activity_prepends_memory_block — configure a fake memory assembler, verify the invocation includes the memory prefix.
   • Add test_dts_activity_memory_failure_does_not_fail_node — configure a raising assembler, verify the node still succeeds.

Gap 3 — Thread mirror writes in DTS activities (P2)

In-process behavior: After each node completes, the orchestrator calls thread_mirror.write_turn() to persist the agent turn (input

• output) to the threads container.

DTS gap: No thread mirror writes.

Fix:

1. Add module-level _thread_mirror reference:

   _thread_mirror: Any | None = None
   
   def configure_thread_mirror(mirror: Any) -> None:
       global _thread_mirror
       _thread_mirror = mirror

2. In execute_node_activity(), after a successful invocation, write the turn (best-effort):

   if _thread_mirror is not None:
       try:
           await _thread_mirror.write_turn(
               run_id=node_input.run_id,
               tenant_id=node_input.tenant_id,
               node_id=node_input.node_id,
               agent_name=node_input.agent_name,
               input_text=call.user_input,
               output_text=result.output_text or "",
           )
       except Exception:
           pass

3. Wire in api.py lifespan.

4. Tests:

   • Add test_dts_activity_writes_thread_turn — configure a fake thread mirror, verify write_turn was called with correct args.

Gap 4 — Service Bus run.completed publish (P1)

In-process behavior: _finalize() calls _sb_publish_completed() to publish run.completed to the per-tenant Service Bus topic {tenant_id}.run.completed. This triggers the Memory Curator.

DTS gap: dag_orchestrator() writes a final Cosmos projection but never publishes to Service Bus. The Memory Curator will never fire for DTS-executed runs.

Fix:

1. Add module-level _sb_publisher reference:

   _sb_publisher: Any | None = None
   
   def configure_sb_publisher(publisher: Any) -> None:
       global _sb_publisher
       _sb_publisher = publisher

2. Add a new DTS activity for publishing the SB event:

   def publish_completion_activity(
       ctx: task.ActivityContext, completion_input: CompletionInput,
   ) -> bool:
       import asyncio
   
       if _sb_publisher is None:
           return False
   
       async def _publish() -> bool:
           await _sb_publisher.publish_run_completed(
               run_id=completion_input.run_id,
               tenant_id=completion_input.tenant_id,
           )
           return True
   
       try:
           return asyncio.run(_publish())
       except Exception:
           logger.warning(
               "SB publish failed for run %s (best-effort)",
               completion_input.run_id,
               exc_info=True,
           )
           return False

3. Add CompletionInput dataclass:

   @dataclass
   class CompletionInput:
       run_id: str
       tenant_id: str

4. In dag_orchestrator(), after the final projection yield, add:

   # Best-effort: publish run.completed to Service Bus
   yield ctx.call_activity(
       publish_completion_activity,
       input=CompletionInput(
           run_id=run_input.run_id,
           tenant_id=run_input.tenant_id,
       ),
   )

5. Wire in api.py lifespan.

6. Tests:

   • Add test_dts_orchestrator_publishes_sb_on_completion — configure a fake SB publisher, run the orchestrator to completion, verify publish was called.

Gap 5 — Working memory writes in DTS activities (P3)

In-process behavior: After each node completes, _apply_result() calls working_memory.write_node_output() to cache the output in Redis for fast downstream lookup.

DTS gap: No working memory writes.

Fix:

1. Add module-level _working_memory reference:

   _working_memory: Any | None = None
   
   def configure_working_memory(writer: Any) -> None:
       global _working_memory
       _working_memory = writer

2. In execute_node_activity(), after success, write to working memory (best-effort):

   if _working_memory is not None and result.output_text:
       try:
           await _working_memory.write_node_output(
               run_id=node_input.run_id,
               node_id=node_input.node_id,
               output_text=result.output_text,
           )
       except Exception:
           pass

3. Wire in api.py lifespan.

4. Tests:

   • Add test_dts_activity_writes_working_memory — configure a fake working memory writer, verify write was called.

Wiring in api.py (applies to all gaps)

The api.py lifespan already configures _invoker and _cosmos_writer for DTS. Extend it to configure the new module-level references:

# In the DTS lifespan block (around line 597-629):
if _dts_enabled(settings):
    from runtime.dts.orchestrator import (
        configure_invoker,
        configure_cosmos_writer,
        configure_memory_assembler,
        configure_thread_mirror,
        configure_sb_publisher,
        configure_working_memory,
    )
    configure_invoker(invoker)
    configure_cosmos_writer(...)
    configure_memory_assembler(app.state.memory_assembler)
    configure_thread_mirror(app.state.thread_mirror)
    configure_sb_publisher(app.state.sb_publisher)
    configure_working_memory(app.state.working_memory)

Also update the TODO(G4) comment at the dispatch site (line ~718) to reflect that parity is now implemented.

File ownership (strictly enforced)

Modify:

• runtime/runtime/dts/orchestrator.py (all 5 gaps)
• runtime/runtime/api.py (wiring only — lifespan + remove TODO)
• runtime/tests/test_dts.py (new tests)

Do NOT touch:

• runtime/runtime/orchestrator/__init__.py (in-process path works)
• runtime/runtime/orchestrator/servicebus_publisher.py
• runtime/runtime/orchestrator/thread_mirror.py
• runtime/runtime/orchestrator/working_memory_writer.py
• runtime/runtime/dts/projection.py
• runtime/runtime/dts/client.py
• runtime/runtime/dts/worker.py
• infra/, synthesizer/, builder/, ui/, data/

Constraints

• No new dependencies.
• All new DTS features are best-effort — failures log warnings but never crash the orchestration or activity.
• The module-level configure pattern is already established by configure_invoker() and configure_cosmos_writer(). Follow the same convention.
• DTS activities are sync functions. Use asyncio.run() to call async code. Reset/close cached async clients before each call to avoid cross-loop errors (the existing pattern).
• Existing tests must continue to pass.
• Run tests after changes:

  cd runtime && .venv/bin/python -m pytest tests/ -x -q --ignore=tests/test_bc_integration.py
  cd runtime && .venv/bin/ruff check runtime/ tests/

  All must pass.