🐛 Decide enum/Literal socket membership once, at assignment

[elinscott]

Problem

A Literal-typed socket's structured-type info dropped its alternatives entirely, so no value assigned to it was ever rejected, whatever it held (see #175 for more details).

The same code path also served Enum-typed sockets, and there a membership check did run, but on two different representations that could disagree with each other.

class Spin(Enum):
    NONE = "none"
    COLLINEAR = "collinear"

class Narrow(Enum):
    """A different enum that happens to repeat one of Spin's values."""
    NONE = "none"

@task.graph()
def narrow_graph(spin: Narrow) -> str:
    return consume_narrow(spin=spin).result

At build time, check_socket_match compared nominal identity — is the source's own enum class a subset of the target's — so passing Spin.NONE (a foreign enum whose value matches Narrow.NONE) was rejected. At run time, the value-reconstruction path compared by value only, so the same foreign-but-matching member was silently accepted once execution reached it. A value that failed one gate could still pass the other, and which gate ran first depended on how the value reached the socket (literal, link, default).

Serialization had a matching gap on the read side: an Enum-typed parameter's serialized form is its bare member value, and nothing rebuilt the member before a task body ran — color.name resolved to str.name instead of raising on a value with no member.

Changes

Read side

Folding in the enum handling from #151 (the earlier enum PR this one supersedes), structured_type_info/coerce_structured_value grow an "enum" kind — serialize to the bare value, rebuild via the class on the way back — so a @task/@task.graph body annotated with an Enum receives the member its signature declares. materialize_graph now also runs the adapter's deserialize over a @task.graph body's resolved inputs before calling the body, so a value round-tripped through an engine-typed wrapper (e.g. aiida-workgraph's orm.Int) arrives as the primitive the signature declares, not a database node.

@task()
def describe(colour: Narrow) -> str:
    return colour.name          # ✅ "NONE" — the body holds the member (main: str.name on the bare value, an AttributeError)

Write side: one decision point

TaskSocket._set_socket_value is the one point every graph shape (direct assignment, link, default, two-hop, namespace) passes through before a value reaches storage, so canonicalization and the allowed-values check now live there instead of being duplicated — the rule lives in one place (canonical_socket_value/value_is_allowed), applied at assignment, defaults, link checks and run-time coercion, by value, never by nominal class. A structured-type socket's default goes through the same check where its spec is built, so a defaulted socket reads the same as an assigned one on both sides of a round trip. For Literal, this is what #175 asked for:

@task()
def paint(colour: Literal["red", "green"]) -> str: ...

paint(colour="red")      # ✅ accepted
paint(colour="green")    # ✅ accepted
paint(colour="blue")     # ❌ ValueError: Invalid value for socket 'paint.inputs.colour'   (main: accepted — Literal carried no alternatives)
paint(colour=17)         # ❌ ValueError, same message                                    (main: accepted)

For an Enum, deciding by value means a foreign enum whose value matches is a member, and a foreign enum whose only match is its name is not:

@task()
def consume_narrow(spin: Narrow) -> str: ...

consume_narrow(spin=Narrow.NONE)   # ✅ member of the declared enum
consume_narrow(spin="none")        # ✅ bare value; the body receives Narrow.NONE, the true member
consume_narrow(spin=Spin.NONE)     # ✅ foreign enum whose VALUE matches: accepted, body receives Narrow.NONE (main: rejected at build, accepted at run)
consume_narrow(spin="collinear")   # ❌ ValueError: no Narrow member has that value
consume_narrow(spin=Spin.COLLINEAR)# ❌ ValueError — matching the NAME of a member is not membership

Accepting Spin.NONE for a Narrow socket is deliberate: after serialization, Narrow.NONE, Spin.NONE and "none" are the same "none", and the socket must take the bare form back from storage — so by value is the only rule a round trip preserves (rejecting at build what run accepted was #175's asymmetry). Static typing is unchanged: mypy still flags spin=Spin.NONE at the call site. The payoff is that a wide enum's member can go straight into a socket declared with a narrower enum of the same values.

Membership by value follows the rule Literal already applies, so an IntEnum no longer takes the values Python merely compares equal to its members:

class Count(IntEnum):
    ONE = 1
    TWO = 2

@task()
def take(n: Count) -> int: ...

take(n=1)      # ✅
take(n=True)   # ❌ ValueError (main: accepted — True == 1)
take(n=1.0)    # ❌ ValueError (main: accepted)

Links

A typed source's declared alternatives must fit the target's, checked at link time; link.py gains check_static_source_value so an untyped source that already holds a literal is checked at build too. A value an untyped task only produces at run cannot be checked at build, so the run-time coercion still raises — the socket names itself in the message, without leaking the TaggedValue wrapper or its uuid.

@task()
def pick_ints() -> Literal[1, 2]: ...

@task()
def passthrough(x: Any) -> Any:
    return x

@task()
def anything() -> Any:
    return "chartreuse"                   # known only when it runs

with Graph() as g:
    a = pick_ints()
    consume_narrow(spin=a.outputs.result) # ❌ TypeError: Socket value range mismatch, at link time

with Graph() as g:
    holder = passthrough(x="chartreuse")  # an untyped socket already holding a literal
    consume_narrow(spin=holder.outputs.x) # ❌ ValueError at build — the static value names no member

with Graph() as g:
    s = anything()
    consume_narrow(spin=s.outputs.result) # builds; ❌ at run: ValueError: Invalid value for socket 'spin'

This closes #175 and answers #176 ("what should a graph body receive") for the case the contract can reach: an annotated socket. A body annotated spin: Narrow receives a value that compares as Narrow.NONE, links onward as it, and lands in the receiving task as the member. What it holds may still be the socket-tagged wrapper — an enum value assigned directly as a graph input stays tagged so the body wires to graph_inputs instead of copying a literal, and TaggedValue(Narrow.NONE) is Narrow.NONE is unavoidably False — but nothing an annotated body needs to do exposes that. An unannotated input is an Any socket: no declared alternatives, so nothing to decide, and whatever the caller passed arrives as-is. Rebuilding a member by hand there was never covered by any contract, and is the one way to see the tag:

@task.graph()
def pick_naive(spin):                          # unannotated input, member rebuilt by hand
    return consume_narrow(spin=Narrow(spin)).result

pick_naive.build(spin="none")
# main:   builds; inside the body, Narrow(spin) == Narrow.NONE is False — Narrow() was handed the TaggedValue, not "none"
# branch: ❌ ValueError: TaggedValue('none', socket=...) is not a valid Narrow

@task.graph()
def pick(spin: Narrow):                        # annotated: nothing to rebuild
    if spin == Narrow.NONE: ...                # ✅ compares as the member
    return consume_narrow(spin=spin).result    # ✅ links; consume_narrow's body receives Narrow.NONE

pick.build(spin="none")                        # ✅
pick.build(spin=Narrow.NONE)                   # ✅

pick_naive — an unannotated graph input passed to Narrow(...) inside the body — is one of the cases #176 tabulates as silently wrong on main: it builds, but the member it constructs compares False against every real Narrow member. On this branch it raises at build. That is deliberate, not a regression: deciding membership by value in one place means a value that names no member fails loudly wherever it appears. The fix is the annotation, as in pick, not a workaround. See the paired aiidateam/aiida-workgraph#800 for the same rule applied on its side of the boundary.

Testing

• tests/test_enum_literal_sockets.py (new, 53 tests) plus 2 in tests/test_engine_local.py: membership decided once — a foreign member with a matching value is accepted at every entry point (literal, link, default, two-hop) and a foreign member with a matching name only is rejected everywhere; Literal of enum members, of strings, of mixed types; requiredness follows the default, not a bare overlay; the run-time message names the socket and never leaks TaggedValue/uuid.
• Negative control (reproduced): reverting just the contract commit while keeping its tests fails 43 of 53 — confirms the tests exercise the single-decision-point change, not something upstream main already did.
• Full suite: 364 passed, no regressions.

[codecov]

Codecov Report

❌ Patch coverage is 94.89194% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.07%. Comparing base (8d85e61) to head (9401e25).

Files with missing lines	Patch %	Lines
tests/test_enum_literal_sockets.py	95.29%	12 Missing ⚠️
src/node_graph/link.py	92.18%	5 Missing ⚠️
src/node_graph/utils/struct_utils.py	93.05%	5 Missing ⚠️
src/node_graph/utils/graph.py	90.00%	2 Missing ⚠️
src/node_graph/socket.py	94.11%	1 Missing ⚠️
src/node_graph/socket_spec.py	98.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #178      +/-   ##
==========================================
+ Coverage   89.68%   90.07%   +0.38%     
==========================================
  Files          81       82       +1     
  Lines        8984     9487     +503     
==========================================
+ Hits         8057     8545     +488     
- Misses        927      942      +15     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.