sequenceDiagram
    participant UI as Widget UI (FE)
    participant Plugin as AnyWidgetPlugin
    participant Model as Frontend Model
    participant Comm as Marimo Comm (BE)
    participant Kernel as Python Kernel

    Note over UI, Kernel: Phase 1: State Capture (During Interaction)

    UI->>Model: model.set("count", 8)
    Model->>Comm: sendUpdate({ method: "echo_update", state: {count: 8} })
    
    rect rgb(240, 240, 240)
        Note over Comm: CHANGED: Handle "echo_update"
        Comm->>Comm: Create ModelUpdate message
        Note right of Comm: Prevents dropping state<br/>on acknowledgement path
    end
    
    Comm-->>Kernel: Broadcast state to all clients

    Note over UI, Kernel: Phase 2: State Hydration (Upon Notebook Reload)

    Kernel->>Comm: Replay stored state
    Comm->>Model: updateAndEmitDiffs({ count: 8 })
    
    Plugin->>Plugin: bind(widgetDef, model)
    Plugin->>UI: render(element, model)
    
    alt Widget Initialization
        UI->>Model: model.on("change:count", listener)
        Note left of UI: Widget defaults to "5"<br/>(Internal initial state)
    end

    rect rgb(240, 240, 240)
        Note over Plugin, Model: NEW: Re-hydration Flow
        Plugin->>Model: getMarimoInternal(model).reemitState()
        
        loop For each key in state
            Model->>UI: emit("change:count", 8)
        end
        
        Model->>UI: emit("change")
    end

    UI->>UI: listener updates UI to "8"

sequenceDiagram
    participant AnyWidget as AnyWidget Frontend
    participant Model as Model (State Manager)
    participant BE as Backend (Comm)
    participant Broadcast as Broadcast System

    Note over AnyWidget,Broadcast: AnyWidget Refresh / Reconnect Flow

    loop On notebook reload / reconnect
        AnyWidget->>Model: get() initial state
        Model-->>AnyWidget: default values (e.g., count=5)
        AnyWidget->>BE: send echo_update (frontend traits)
        BE->>Broadcast: broadcast Notification (ModelUpdate)
        Broadcast-->>AnyWidget: state event
        Note over AnyWidget,Model: state may arrive before render listener is attached
    end

    alt Late-attached render listener
        AnyWidget->>Model: reemitState()
        Model->>Model: #reemitState: emit change:key events for each current value
        Model--)AnyWidget: change:count (e.g., count=8)
        Note over AnyWidget: Listener receives hydrated state even if late
    else Normal path (listener already attached during connect)
        AnyWidget->>Model: reemitState()
        Model->>Model: #reemitState
        Model--)AnyWidget: change events (duplicate but harmless)
    end

    Note over AnyWidget,Broadcast: echo_update on BE side

    BE->>BE: _create_model_message(method="echo_update")
    alt echo_update method
        BE->>BE: Convert to ModelUpdate (preserves state)
        BE->>Broadcast: broadcast with state
    else Unknown method
        BE->>BE: Return None (skip)
    end