Merge pull request #294 from fireflyframework/feat/workflows-streaming-generics

feat(workflows): token-level streaming and typed generics

Author: ancongui

CHANGELOG.md

@@ -7,6 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 Copyright 2026 Firefly Software Foundation. Licensed under the Apache License 2.0.
 
+## [26.06.4] - 2026-06-21
+
+Dynamic Workflows — the final SOTA wave: token-level streaming and end-to-end
+static typing of the DSL.
+
+### Added
+
+- **Streaming** — `stream(prompt, ...)` is an async context manager that streams one
+  sub-agent's output token-by-token: iterate `handle.text()` for deltas, then read
+  `handle.output` for the full result after the block. It honours the budget,
+  concurrency gate, journal (a resumed call yields its cached output once) and cost
+  accounting exactly like `agent()`. Backed by a `StreamingAgentRunner` protocol that
+  `DefaultAgentRunner` implements via pydantic-ai's `run_stream`; a non-streaming
+  runner raises `WorkflowError`. Streamed calls emit `agent.start` / `agent.end` with
+  `stream: True`. New exports: `stream`, `StreamHandle`, `StreamingAgentRunner`.
+- **Typed generics** — `agent(output_type=T)` is now typed to return `T` (via
+  `@overload`) instead of `Any`, and `@workflow` produces a `Workflow[OutputT]`
+  inferred from the function's return annotation, so `await my_workflow(args)` is
+  statically typed end-to-end with no casts.
+
 ## [26.06.3] - 2026-06-20
 
 Dynamic Workflows — the durable-composition wave: sub-workflows, durable resume,

docs/workflows.md

@@ -77,6 +77,7 @@ or by name through `run_workflow("deep_research", args)`.
 | `agent` | `await agent(prompt, *, label, model, output_type, instructions, deps, tools, toolsets)` | Run one isolated sub-agent; returns its `output` (a `str` or a validated `output_type`). A sub-agent can use `tools=`/`toolsets=` (e.g. `ToolKit.as_toolset()` or an MCP server) just like a top-level agent. Honours the budget, the concurrency gate, and the resume journal. |
 | `parallel` | `await parallel(thunks)` | **Barrier.** Run zero-arg async thunks concurrently; a thunk that raises resolves to `None` (the call never propagates). Returns a list aligned to `thunks`. |
 | `pipeline` | `await pipeline(items, *stages)` | **No inter-stage barrier.** Each item flows through every stage independently (item A can be in stage 3 while B is in stage 1). Each stage receives `(prev, item, index)` — declare only the params you need. A stage that raises drops *that* item to `None`. |
+| `stream` | `async with stream(prompt, ...) as s:` | Stream one sub-agent token-by-token: iterate `s.text()` for deltas; `s.output` holds the full output after the block. Same budget/journal/cost accounting as `agent`. Requires a streaming runner. See [Streaming](#streaming). |
 | `phase` | `with phase("title"):` | Group enclosed work for telemetry (`phase.start` / `phase.end` events). |
 | `log` | `log("message")` | Emit a narrator line to the run's event handler. |
 
@@ -85,6 +86,30 @@ glue between agent calls — dedup, rank, filter, branch — is ordinary
 deterministic Python. Reach for an `agent()` only when you genuinely need a
 model.
 
+### Type safety
+
+The DSL is statically typed end-to-end. `agent(output_type=T)` is typed to
+return `T` (not `Any`), and `@workflow` produces a `Workflow[OutputT]` inferred
+from the function's return annotation — so the awaited result is typed too:
+
+```python
+from pydantic import BaseModel
+from fireflyframework_agentic.workflows import workflow, agent
+
+class Report(BaseModel):
+    summary: str
+    sources: list[str]
+
+@workflow(name="research")
+async def research(args, ctx) -> Report:
+    return await agent("write the report", output_type=Report)   # typed Report
+
+report: Report = await research(args)   # research is Workflow[Report]; result is Report
+```
+
+Without `output_type`, `agent()` returns `Any` (the raw string output). Pyright
+infers all of the above with no casts.
+
 ---
 
 ## Budgets
@@ -177,6 +202,29 @@ await deep_research(args, runner=MyFakeRunner())
 
 ---
 
+## Streaming
+
+Stream a single sub-agent's output token-by-token with the `stream()` context
+manager (requires a streaming-capable runner — `DefaultAgentRunner` is). It
+honours the budget, concurrency gate, journal and cost accounting exactly like
+`agent()`:
+
+```python
+from fireflyframework_agentic.workflows import workflow, stream
+
+@workflow(name="live_report")
+async def live_report(args, ctx):
+    async with stream("write a cited report", model=args.model) as s:
+        async for delta in s.text():
+            print(delta, end="", flush=True)   # surface tokens as they arrive
+    return s.output                            # the full text, available after the block
+```
+
+On resume, a cached `stream()` call yields its full output once (the model is not
+re-invoked). A runner that doesn't support streaming raises `WorkflowError`.
+
+---
+
 ## Model routing & cost optimization
 
 The cheapest model that can do the job should do the job. Two complementary tools
@@ -326,7 +374,8 @@ process restart.
 Pass `events=callable` to a run to receive structured events:
 `workflow.start` / `workflow.end` (with `agents`, `tokens`, `cost_usd`),
 `phase.start` / `phase.end`, `agent.start` / `agent.end` (with `label`, `phase`,
-`seq`, `tokens`, `cost_usd`), `route.select` / `route.escalate`, `cascade.tier`,
+`seq`, `tokens`, `cost_usd`; streamed calls also carry `stream: True`),
+`route.select` / `route.escalate`, `cascade.tier`,
 `subworkflow.start` / `subworkflow.end`, `human.pause`, and `log`. Wire this into
 the [observability](observability.md) layer for live per-phase token/cost/agent/time
 counters.
@@ -342,6 +391,8 @@ counters.
 | `run_workflow(name, args, **opts)` | Look up a registered workflow by name and run it. |
 | `subworkflow(name_or_wf, args)` | Run another workflow inline, inheriting the parent's budget/journal/runner. |
 | `agent` / `parallel` / `pipeline` / `phase` / `log` | The DSL primitives. |
+| `stream(prompt, ...)` | Async context manager → a `StreamHandle` (`.text()` deltas, `.output`); streams one sub-agent. |
+| `StreamHandle` / `StreamingAgentRunner` | The streaming result handle and the runner protocol a runner implements to support `stream`. |
 | `human(prompt)` | Pause for external input (raises `WorkflowInterrupt`; resume via the same journal). |
 | `map_agents(items, fn, *, strict=False)` | Run `fn(item)` per item concurrently — sugar over `parallel` (no late-binding lambda). |
 | `WorkflowBudget` | Concurrency / agent-count / token / **USD cost** / **wall-clock** ceilings. |

fireflyframework_agentic/workflows/__init__.py

@@ -53,6 +53,7 @@ async def deep_research(args, ctx):
     parallel,
     phase,
     pipeline,
+    stream,
 )
 from fireflyframework_agentic.workflows.registry import (
     Workflow,
@@ -73,6 +74,8 @@ async def deep_research(args, ctx):
     AgentCall,
     AgentRunner,
     DefaultAgentRunner,
+    StreamHandle,
+    StreamingAgentRunner,
 )
 from fireflyframework_agentic.workflows.verify import (
     CascadeResult,
@@ -95,6 +98,8 @@ async def deep_research(args, ctx):
     "JournalBackend",
     "ModelSelectionStrategy",
     "SmartRoutingRunner",
+    "StreamHandle",
+    "StreamingAgentRunner",
     "Verdict",
     "Workflow",
     "WorkflowBudget",
@@ -115,6 +120,7 @@ async def deep_research(args, ctx):
     "pipeline",
     "price_model",
     "run_workflow",
+    "stream",
     "subworkflow",
     "workflow",
     "workflow_registry",

fireflyframework_agentic/workflows/primitives.py

@@ -35,23 +35,54 @@
 import contextlib
 import inspect
 import logging
-from collections.abc import Awaitable, Callable, Iterable, Iterator
-from typing import Any
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator
+from typing import Any, TypeVar, overload
 
 from fireflyframework_agentic.exceptions import (
     WorkflowBudgetError,
     WorkflowContextError,
+    WorkflowError,
     WorkflowInterrupt,
 )
 from fireflyframework_agentic.workflows.context import current_workflow
+from fireflyframework_agentic.workflows.runner import AgentCall, StreamingAgentRunner
 
 logger = logging.getLogger(__name__)
+_T = TypeVar("_T")
 
 # Structural / kill-switch / control-flow signals must abort the run, not be
 # swallowed to ``None`` by a parallel branch or a pipeline stage.
 _NEVER_SWALLOW = (WorkflowBudgetError, WorkflowContextError, WorkflowInterrupt)
 
 
+@overload
+async def agent(
+    prompt: Any,
+    *,
+    label: str | None = ...,
+    model: Any | None = ...,
+    output_type: type[_T],
+    instructions: str | None = ...,
+    deps: Any = ...,
+    tools: Any | None = ...,
+    toolsets: Any | None = ...,
+) -> _T: ...
+
+
+@overload
+async def agent(
+    prompt: Any,
+    *,
+    label: str | None = ...,
+    model: Any | None = ...,
+    output_type: None = ...,
+    instructions: str | None = ...,
+    deps: Any = ...,
+    tools: Any | None = ...,
+    toolsets: Any | None = ...,
+) -> Any: ...
+
+
 async def agent(
     prompt: Any,
     *,
@@ -142,6 +173,88 @@ async def human(prompt: str, *, label: str | None = None) -> Any:
     raise WorkflowInterrupt(prompt, seq=seq, journal=ctx.journal)
 
 
+class _CachedStream:
+    """Stream handle for a journal-cached call: yields the full output once."""
+
+    def __init__(self, output: Any) -> None:
+        self._output = output
+        self.call = AgentCall(output=output)
+
+    async def text(self) -> AsyncIterator[str]:
+        yield str(self._output)
+
+    @property
+    def output(self) -> Any:
+        return self._output
+
+
+@contextlib.asynccontextmanager
+async def stream(
+    prompt: Any,
+    *,
+    label: str | None = None,
+    model: Any | None = None,
+    output_type: Any | None = None,
+    instructions: str | None = None,
+    deps: Any = None,
+    tools: Any | None = None,
+    toolsets: Any | None = None,
+) -> AsyncIterator[Any]:
+    """Stream one sub-agent's output token-by-token, as a context manager.
+
+    Iterate ``handle.text()`` for deltas; after the block, ``handle.output`` holds
+    the full output. Honours the budget, concurrency gate, journal (a resumed call
+    yields its cached output once) and cost accounting, exactly like ``agent()``.
+    Requires the run's runner to support streaming (``DefaultAgentRunner`` does).
+
+    Example::
+
+        async with stream("write a cited report", model=args.model) as s:
+            async for delta in s.text():
+                print(delta, end="", flush=True)
+        report = s.output
+    """
+    ctx = current_workflow()
+    seq = ctx.next_call_seq()
+    if ctx.journal.has(seq):
+        yield _CachedStream(ctx.journal.get(seq))
+        return
+
+    runner = ctx.runner
+    if not isinstance(runner, StreamingAgentRunner):
+        raise WorkflowError(f"the active runner {type(runner).__name__} does not support streaming")
+
+    ctx.reserve_agent()
+    display = label or (prompt[:40] if isinstance(prompt, str) else f"stream#{seq}")
+    ctx.emit("agent.start", {"label": display, "phase": ctx.current_phase, "seq": seq, "stream": True})
+    async with ctx.semaphore:
+        async with runner.run_stream(
+            prompt,
+            model=model,
+            output_type=output_type,
+            instructions=instructions,
+            deps=deps,
+            tools=tools,
+            toolsets=toolsets,
+        ) as handle:
+            yield handle
+        call = handle.call if handle.call is not None else AgentCall(output=handle.output)
+    ctx.record_tokens(call.tokens)
+    ctx.record_cost_usd(call.cost_usd)
+    ctx.journal.record(seq, call.output)
+    ctx.emit(
+        "agent.end",
+        {
+            "label": display,
+            "phase": ctx.current_phase,
+            "seq": seq,
+            "tokens": call.tokens,
+            "cost_usd": call.cost_usd,
+            "stream": True,
+        },
+    )
+
+
 async def parallel(thunks: Iterable[Callable[[], Awaitable[Any]]]) -> list[Any]:
     """Run zero-arg async ``thunks`` concurrently and return their results.
 

fireflyframework_agentic/workflows/registry.py

@@ -28,7 +28,8 @@
 import logging
 import threading
 import uuid
-from typing import Any
+from collections.abc import Awaitable, Callable
+from typing import Any, Generic, TypeVar
 
 from fireflyframework_agentic.exceptions import WorkflowBudgetError, WorkflowNotFoundError
 from fireflyframework_agentic.workflows.context import (
@@ -43,6 +44,7 @@
 logger = logging.getLogger(__name__)
 
 EventHandler = Any  # Callable[[str, dict], None]
+OutputT = TypeVar("OutputT")
 
 
 def _wants_context(fn: Any) -> bool:
@@ -58,13 +60,13 @@ def _wants_context(fn: Any) -> bool:
     return len(params) >= 2
 
 
-class Workflow:
-    """A registered, runnable workflow."""
+class Workflow(Generic[OutputT]):
+    """A registered, runnable workflow (generic over its return type)."""
 
     def __init__(
         self,
         name: str,
-        fn: Any,
+        fn: Callable[..., Awaitable[OutputT]],
         *,
         args_schema: Any | None = None,
         description: str = "",
@@ -83,7 +85,7 @@ async def run(
         journal: Journal | None = None,
         events: EventHandler | None = None,
         run_id: str | None = None,
-    ) -> Any:
+    ) -> OutputT:
         """Execute the workflow, returning whatever its body returns.
 
         Pass the same ``journal`` from a prior run to resume: completed agent
@@ -127,7 +129,7 @@ async def run(
         finally:
             _current.reset(token)
 
-    async def __call__(self, args: Any = None, **kwargs: Any) -> Any:
+    async def __call__(self, args: Any = None, **kwargs: Any) -> OutputT:
         return await self.run(args, **kwargs)
 
 
@@ -177,20 +179,23 @@ def workflow(
     args_schema: Any | None = None,
     description: str = "",
     register: bool = True,
-) -> Any:
+) -> Callable[[Callable[..., Awaitable[OutputT]]], Workflow[OutputT]]:
     """Decorator: turn an async function into a registered :class:`Workflow`.
 
+    The returned :class:`Workflow` is generic over the function's return type, so
+    ``await my_workflow(args)`` is typed as that return type rather than ``Any``.
+
     Example::
 
         @workflow(name="deep_research", args_schema=ResearchArgs)
-        async def deep_research(args, ctx):
+        async def deep_research(args, ctx) -> Report:
             with phase("search"):
                 hits = await parallel([lambda q=q: agent(f"search: {q}") for q in args.queries])
-            return await agent("synthesize", deps=hits)
+            return await agent("synthesize", deps=hits, output_type=Report)
     """
 
-    def decorator(fn: Any) -> Workflow:
-        wf_name = name or fn.__name__
+    def decorator(fn: Callable[..., Awaitable[OutputT]]) -> Workflow[OutputT]:
+        wf_name = name or getattr(fn, "__name__", "workflow")
         wf = Workflow(wf_name, fn, args_schema=args_schema, description=description)
         if register:
             workflow_registry.register(wf_name, wf)