Retry, STOP & rollback.

$ git clone https://github.com/Cosmonapse/cosmonapse-examples
$ pip install cosmonapse

$ cd cosmonapse-examples/13-retry
$ python brain.py all          # the three scenarios in order
$ python brain.py retry
$ python brain.py give-up
$ python brain.py rollback

1. retry survives a stalled stage
   attempt 1 stuck (TimeoutError) -> STOP + re-dispatch on a fresh trace
   answer on attempt 2: answer to 'what is a Dendrite?'

2. retry gives up
   gave up after 2 attempts, each STOPped (2 invocations)

3. roll back a half-finished ingest
   index size before ingest: 0
   ingest result: ERROR (ingest of 'doc' crashed after 3 chunks)
   index size after crash: 3  (partial write)
   stop_trace(rollback=True): 3 inverse ops replayed
   index size after rollback: 0  (clean)

13-retry/
  config.py              appends ../11-rag to sys.path
  index.py               VECTORS - 11-rag's VectorEngram
  neurons/ingester.py    imprints chunks; crashes after fail_after on request
  neurons/generator.py   stalls for the first stall_first calls per question
  receptors/terminal.py  retry, give-up, rollback, all
  brain.py               index-node + ingester-node + generator-node + terminal-node

import asyncio
from collections import Counter

from cosmonapse import Axon

CALLS: Counter = Counter()          # question -> invocations, for the printout


async def generator(input, context):
    question = input["question"]
    CALLS[question] += 1
    n = CALLS[question]
    if n <= input.get("stall_first", 0):
        await asyncio.sleep(input.get("stall_s", 2.0))   # past the retry timeout
    return {"answer": f"answer to {question!r}", "attempt": n}


AXON = Axon(neuron_id="generator", neuron_fn=generator, capabilities=["generate"])

async def ingest_neuron(input, context, *, imprint):
    """{"doc_id", "text", "fail_after"?} -> {"doc_id", "chunks"}.

    Imprints each chunk into the vector index under the TASK's trace. When
    ``fail_after`` is set the neuron raises partway through, leaving a partial
    write the orchestrator can roll back.
    """
    doc_id = input["doc_id"]
    chunks = chunk_text(input["text"])
    fail_after = input.get("fail_after")
    written = 0
    for i, chunk in enumerate(chunks):
        if fail_after is not None and i >= fail_after:
            raise RuntimeError(
                f"ingest of {doc_id!r} crashed after {written} chunks"
            )
        await imprint(
            "vectors", op="upsert",
            entry={"doc_id": doc_id, "chunk_index": i, "text": chunk,
                   "embedding": fake_embed(chunk)},
            merge_key=f"{doc_id}:{i}", await_ack=True, deadline_ms=2000,
        )
        written += 1
    return {"doc_id": doc_id, "chunks": written}


AXON = Axon(
    neuron_id="ingester", neuron_fn=ingest_neuron,
    capabilities=["ingest"],
    engrams=[EngramBinding(name="vectors", directed_id="retry-vectors")],
)

@RECEPTOR.command("retry", local=True, help="retry survives a stalled stage")
async def retry():
    lines = []

    def on_retry(attempt, outcome):
        # attempt is 0-based: the index of the attempt that just failed
        lines.append(f"attempt {attempt + 1} stuck ({type(outcome).__name__}) "
                     f"-> STOP + re-dispatch on a fresh trace")

    question = "what is a Dendrite?"
    sig = await RECEPTOR.dendrite.run_with_retry(
        neuron="generator",
        input={"question": question, "stall_first": 1},
        retry=RetryStrategy(max_attempts=3, timeout_s=0.5, on_retry=on_retry),
    )
    out = sig.payload["output"]
    lines.append(f"answer on attempt {out['attempt']}: {out['answer']}")
    return "\n".join(lines)


@RECEPTOR.command("give-up", local=True, help="retry exhausts its attempts")
async def give_up():
    question = "anything"
    try:
        await RECEPTOR.dendrite.run_with_retry(
            neuron="generator",
            input={"question": question, "stall_first": 99},
            retry=RetryStrategy(max_attempts=2, timeout_s=0.4),
        )
    except asyncio.TimeoutError:
        return f"gave up after 2 attempts, each STOPped ({CALLS[question]} invocations)"
    return "unexpected: an attempt answered"

@RECEPTOR.command("rollback", local=True, help="undo a half-finished ingest")
async def rollback():
    orch, lines = RECEPTOR.dendrite, []
    lines.append(f"index size before ingest: {VECTORS.size}")
    tid = new_trace_id()
    reply = await orch.dispatch_and_wait(
        neuron="ingester",
        input={"doc_id": "doc", "text": DOC, "fail_after": 3},
        trace_id=tid, timeout_s=10.0,
    )
    lines.append(f"ingest result: {reply.type.value} "
                 f"({reply.payload.get('message', 'ok')})")
    lines.append(f"index size after crash: {VECTORS.size}  (partial write)")
    acks = await orch.stop_trace(tid, rollback=True, collect_acks=True, timeout_s=0.5)
    replayed = sum(a.payload.get("compensated", 0) for a in acks)
    lines.append(f"stop_trace(rollback=True): {replayed} inverse ops replayed")
    lines.append(f"index size after rollback: {VECTORS.size}  (clean)")
    return "\n".join(lines)

from engram.vector_engram import VectorEngram      # from 11-rag

VECTORS = VectorEngram(engram_id="retry-vectors", engram_kind="semantic")

import asyncio
import contextlib
import signal

from cosmonapse import Dendrite, MemorySynapse, connect_synapse, run_brain

from config import NAMESPACE, SYNAPSE_URL
from index import VECTORS
from neurons import generator, ingester
from receptors import terminal


def build_index(synapse) -> Dendrite:
    node = Dendrite(synapse=synapse, namespace=NAMESPACE,
                    dendrite_id="index-node", role="worker")
    node.attach_engram(VECTORS)
    return node


def build_neuron(synapse, axon) -> Dendrite:
    node = Dendrite(synapse=synapse, namespace=NAMESPACE,
                    dendrite_id=f"{axon.neuron_id}-node", role="worker")
    node.attach_axon(axon)
    return node


def build_terminal(synapse) -> Dendrite:
    node = Dendrite(synapse=synapse, namespace=NAMESPACE,
                    dendrite_id="terminal-node", heartbeat_s=0)
    node.attach_receptor(terminal.RECEPTOR)
    return node

# terminal 1 - a dev synapse (TCP, single host)
$ cosmo synapse start memory --namespace=retry-demo

# terminal 2 - the same brain, over that synapse
$ SYNAPSE_URL=cosmo://127.0.0.1:7070 python brain.py all

# terminal 3 - Prism, the live view (http://127.0.0.1:7071)
$ cosmo prism --url=cosmo://127.0.0.1:7070 -n retry-demo

# production transports - only the URL changes
$ SYNAPSE_URL=nats://127.0.0.1:4222 python brain.py all
$ SYNAPSE_URL=kafka://127.0.0.1:9092 python brain.py all