feat(flow): let a chat flow declare its own state shape (#7061)

* feat(flow): let a chat flow declare its own state shape

A conversational declaration could only use `state: {type: pydantic, ref: ...}`
pointing at a `ConversationState` subclass. Every other shape loaded clean and
then died on the first turn -- inline `json_schema` and a non-subclass ref with
`AttributeError: 'StateWithId' object has no attribute 'messages'`, and
`type: dict` with `AttributeError: 'dict' object has no attribute 'id'`.

`Flow._compose_extension_state_model` is a new runtime extension seam -- the
seventh alongside the existing six -- applied to the model built from `state:`
before the engine wraps it for `id`. The conversational mixin uses it to add
the chat fields to whatever the declaration asked for, so declared fields and
defaults survive; a model that already extends `ConversationState` is returned
untouched, so today's supported shape is a no-op.

`dict` and `unknown` state cannot carry those fields at all, so the default
extension state supplies the real shape (seeded from the declared defaults
where they fit) rather than forbidding it. Raising instead would break
construction, and `Flow[dict]` with `conversational = True` constructs today --
`crewai flow plot` and definition-only consumers would stop working.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix(flow): keep declared defaults and cover an unbuildable state model

Two review follow-ups on the declared-state work.

A `dict` state's defaults are arbitrary keys, and the fallback kept only the
ones matching `ConversationState`, so `{"type": "dict", "default": {"topic":
"ai"}}` lost `topic` before the first turn and an action reading `state.topic`
would fail. They are carried as extras now.

And a declared `pydantic`/`json_schema` state whose model cannot be built --
a bad ref, an invalid schema -- fell through to a plain dict with none of the
chat fields, so the turn died on `state.id` instead. The engine now re-asks the
extension in that case, as if nothing had been declared.

Found by Cursor and CodeRabbit on #7061.

* refactor(flows): drop the unreachable extension-state fallback

The _initial_state_t branch sat after an unconditional return. The
state_definition is None case and _conversation_state_with_defaults now cover
every path that used to reach it.

Found by Cursor on #7061.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
João Moura
2026-08-24 13:39:46 -03:00
committed by GitHub
parent 9e9a8577be
commit f68fd9e850
3 changed files with 241 additions and 18 deletions

View File

@@ -23,9 +23,9 @@ from contextlib import contextmanager
from enum import Enum
import json
import logging
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, cast
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
from pydantic import BaseModel, Field, create_model
from pydantic import BaseModel, ConfigDict, Field, create_model
from crewai.events.event_bus import crewai_event_bus
from crewai.events.types.flow_events import (
@@ -140,6 +140,37 @@ def _config_from_definition(
)
def _conversation_state_with_defaults(
defaults: dict[str, Any] | None,
) -> ConversationState:
"""Build the chat state, keeping declared defaults the shape cannot hold.
A ``dict`` state's defaults are arbitrary keys, and dropping the ones that
are not conversational fields would break an action reading
``state.topic``. They are carried as extras instead.
"""
if not defaults:
return ConversationState()
known = {
name: value
for name, value in defaults.items()
if name in ConversationState.model_fields
}
extra = {
name: value
for name, value in defaults.items()
if name not in ConversationState.model_fields
}
if not extra:
return ConversationState(**known)
class ConversationStateWithDeclaredDefaults(ConversationState):
model_config = ConfigDict(extra="allow")
return ConversationStateWithDeclaredDefaults(**known, **extra)
def _builtin_method(handler: Callable[..., Any], **roles: Any) -> FlowMethodDefinition:
"""One built-in conversational method, as the code ref the DSL already emits."""
return FlowMethodDefinition(do=_method_action(handler), **roles)
@@ -997,23 +1028,53 @@ class _ConversationalMixin:
update={"methods": {**definition.methods, **missing}}
)
def _create_default_extension_state(self) -> ConversationState | None:
"""Supply ``ConversationState`` only when nothing else declares state.
def _compose_extension_state_model(
self, model_class: type[BaseModel]
) -> type[BaseModel]:
"""Add the conversational fields to a declared state model.
A declared ``state:`` block always wins. This hook is consulted before
``_create_definition_state``, so returning a state here would discard
every field the declaration asked for.
A turn reads ``messages``, ``current_user_message``, ``last_intent``,
``ended``, ``events`` and ``agent_threads`` off state, so a declared
model that lacks them fails on the first turn. Composing rather than
replacing keeps every field the declaration asked for. A model that
already extends ``ConversationState`` is returned untouched.
"""
if not self._is_conversational_enabled():
return model_class
if issubclass(model_class, ConversationState):
return model_class
class ConversationalState(ConversationState, model_class): # type: ignore[misc, valid-type]
pass
return ConversationalState
def _create_default_extension_state(
self, *, ignore_declared_state: bool = False
) -> ConversationState | None:
"""Supply ``ConversationState`` when the declaration cannot carry it.
A declared ``pydantic`` or ``json_schema`` state is composed with the
conversational fields instead (see
``_compose_extension_state_model``), so it keeps everything it asked
for. ``dict`` and ``unknown`` state cannot carry them at all, so this
supplies the real shape rather than letting the turn die on
``state.id`` -- seeded from the declared defaults where they fit.
"""
if not self._is_conversational_enabled():
return None
if self._conversation_flow_definition().state is not None:
return None
initial_state_t = getattr(self, "_initial_state_t", None)
if not hasattr(self, "_initial_state_t") or isinstance(
initial_state_t, TypeVar
):
state_definition = self._conversation_flow_definition().state
if state_definition is None:
return ConversationState()
return None
if not ignore_declared_state and state_definition.type in (
"pydantic",
"json_schema",
):
# Composed with the declared model instead; see
# ``_compose_extension_state_model``.
return None
return _conversation_state_with_defaults(state_definition.default)
def _should_apply_pending_kickoff_context(self) -> bool:
return (

View File

@@ -185,6 +185,7 @@ def _condition_satisfied(condition: FlowDefinitionCondition, events: set[str]) -
def _build_definition_state_model(
state_definition: FlowStateDefinition,
compose: Callable[[type[BaseModel]], type[BaseModel]] | None = None,
) -> BaseModel | None:
kwargs = dict(state_definition.default or {})
@@ -216,6 +217,9 @@ def _build_definition_state_model(
if model_class is None:
return None
if compose is not None:
model_class = compose(model_class)
if not issubclass(model_class, FlowState):
class StateWithId(FlowState, model_class): # type: ignore[misc, valid-type]
@@ -458,10 +462,28 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
"""
return definition
def _create_default_extension_state(self) -> Any | None:
"""Return a default state supplied by an optional runtime extension."""
def _create_default_extension_state(
self, *, ignore_declared_state: bool = False
) -> Any | None:
"""Return a default state supplied by an optional runtime extension.
``ignore_declared_state`` is set when a declared ``state:`` block named
a model that could not be built, so the extension is asked again as if
nothing had been declared.
"""
return None
def _compose_extension_state_model(
self, model_class: type[BaseModel]
) -> type[BaseModel]:
"""Let an optional runtime extension add bases to a declared state model.
Applied to the model built from ``state:`` before the engine wraps it
for its ``id`` field, so an extension can require its own fields
alongside whatever the declaration asked for.
"""
return model_class
def _should_apply_pending_kickoff_context(self) -> bool:
"""Whether an optional runtime extension has pending kickoff context."""
return False
@@ -1721,7 +1743,9 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
if state_definition is None:
return {"id": str(uuid4())}
if state_definition.type in ("pydantic", "json_schema"):
state = _build_definition_state_model(state_definition)
state = _build_definition_state_model(
state_definition, compose=self._compose_extension_state_model
)
if state is not None:
return state
logger.error(
@@ -1730,6 +1754,11 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
self._definition.name,
state_definition.type,
)
extension_state: dict[str, Any] | BaseModel | None = (
self._create_default_extension_state(ignore_declared_state=True)
)
if extension_state is not None:
return extension_state
elif state_definition.type == "unknown":
logger.warning(
"Flow %r declares state of unknown type; falling back to dict state",

View File

@@ -2420,6 +2420,140 @@ class TestDeclarativeConversationalFlow:
assert "cannot carry a model class" in caplog.text
class DeclaredSchemaChatState(ConversationState):
"""A ConversationState subclass, i.e. the already-supported state shape."""
ticket_id: str | None = None
class TestDeclaredConversationalState:
"""Any declared state shape works for a chat flow, keeping its own fields."""
@staticmethod
def _flow(state: dict[str, Any] | None) -> Flow[Any]:
declaration = _conversational_declaration(conversational={})
if state is None:
declaration.pop("state", None)
else:
declaration["state"] = state
return Flow.from_declaration(contents=declaration)
CONVERSATIONAL_FIELDS = (
"messages",
"current_user_message",
"last_intent",
"ended",
"events",
"agent_threads",
)
def test_inline_json_schema_state_gains_the_conversational_fields(self) -> None:
flow = self._flow(
{
"type": "json_schema",
"json_schema": {
"type": "object",
"properties": {
"ticket_id": {"type": "string"},
"turns": {"type": "integer"},
},
},
"default": {"ticket_id": "T-1"},
}
)
fields = type(flow.state).model_fields
for name in self.CONVERSATIONAL_FIELDS:
assert name in fields, name
assert "ticket_id" in fields and "turns" in fields
assert flow.state.ticket_id == "T-1"
assert flow.state.id
def test_a_pydantic_ref_that_is_not_a_conversation_state_is_composed(self) -> None:
flow = self._flow({"type": "pydantic", "ref": "crewai.flow:ChatState"})
fields = type(flow.state).model_fields
for name in self.CONVERSATIONAL_FIELDS:
assert name in fields, name
assert "session_ready" in fields
def test_a_conversation_state_subclass_is_not_composed_twice(self) -> None:
flow = self._flow(
{"type": "pydantic", "ref": f"{__name__}:DeclaredSchemaChatState"}
)
mro = [cls.__name__ for cls in type(flow.state).__mro__]
assert mro.count("ConversationState") == 1
assert "ticket_id" in type(flow.state).model_fields
def test_dict_state_is_given_the_real_shape(self) -> None:
"""A dict cannot carry the fields, so supply them rather than die."""
flow = self._flow({"type": "dict", "default": {"last_intent": "order"}})
assert isinstance(flow.state, ConversationState)
assert flow.state.last_intent == "order"
def test_dict_state_keeps_declared_defaults_it_does_not_own(self) -> None:
"""Dropping them would break an action reading `state.topic`."""
flow = self._flow({"type": "dict", "default": {"topic": "ai", "limit": 3}})
assert flow.state.topic == "ai"
assert flow.state.limit == 3
assert isinstance(flow.state, ConversationState)
def test_unknown_state_is_treated_like_dict(self) -> None:
flow = self._flow({"type": "unknown", "ref": "x:Y", "default": {"topic": "ai"}})
assert flow.state.topic == "ai"
assert isinstance(flow.state, ConversationState)
def test_an_unbuildable_declared_model_still_gets_the_chat_shape(self) -> None:
"""A bad ref fell back to a plain dict, so the turn died on `state.id`."""
flow = self._flow({"type": "pydantic", "ref": "no.such.module:Nope"})
assert isinstance(flow.state, ConversationState)
assert flow.state.id
def test_declared_state_still_runs_a_turn(self) -> None:
flow = self._flow(
{
"type": "json_schema",
"json_schema": {
"type": "object",
"properties": {"ticket_id": {"type": "string"}},
},
}
)
flow._conversation_config.llm = _ScriptedLLM(["Hello."])
assert flow.handle_turn("hi") == "Hello."
assert [m.role for m in flow.state.messages] == ["user", "assistant"]
def test_non_conversational_declaration_keeps_its_plain_state(self) -> None:
flow = Flow.from_declaration(
contents={
"schema": "crewai.flow/v1",
"name": "Plain",
"state": {
"type": "json_schema",
"json_schema": {
"type": "object",
"properties": {"topic": {"type": "string"}},
},
"default": {"topic": "ai"},
},
"methods": {
"begin": {
"do": {"call": "expression", "expr": "state.topic"},
"start": True,
}
},
}
)
assert "messages" not in type(flow.state).model_fields
assert flow.state.topic == "ai"
class TestDeclaredConversationalLLM:
"""A conversational block accepts the shapes a crew agent's `llm` accepts."""
@@ -2525,7 +2659,6 @@ class TestDeclaredConversationalLLM:
assert declared in seen
assert [m.role for m in flow.state.messages][0] == "user"
class TestRoutingArtefactLabels:
"""A route label echoed by a handler is a routing artefact, not a reply."""