Protocol Messages Refactor

Status: Planning — not yet implemented
Motivation: Eliminate implicit payload contracts scattered across participants. Messages become self-describing objects; participants no longer know string key names.

1. Problem

The bus provides transport-level loose coupling, but payload structure is an implicit contract encoded in string literals at every publish and consume site:

# Publisher (generator_agent.py)
{"chunk": chunk.message.thinking, "session_id": session_id, "query_id": query_id}

# Consumer (ws_bridge.py — has to know the same keys)
payload.get("chunk", "")
payload.get("session_id")

Rename a key anywhere → silent breakage. Add a field → no way to know all consumers. The message_type string ("thinking", "response", "critique") is also redundant with the subject — set manually at every publish site.

2. Design

2.1 Message class anatomy

# src/local/protocol/messages.py

from __future__ import annotations
from dataclasses import dataclass, field, asdict
from typing import ClassVar
from local.protocol.subjects import GENERATION_THINKING

@dataclass
class GenerationThinking:
    subject:        ClassVar[str] = GENERATION_THINKING
    message_type:   ClassVar[str] = "thinking"          # derived from class — no longer passed manually

    chunk:      str
    session_id: str
    query_id:   str

    # --- bus boundary ---

    def to_envelope(self, sender_id: str, correlation_id: str) -> MessageEnvelope:
        return MessageEnvelope.create(
            message_type=self.message_type,
            subject=self.subject,
            sender_id=sender_id,
            payload=asdict(self),
            correlation_id=correlation_id,
            metadata={"session_id": self.session_id},
        )

    @classmethod
    def from_envelope(cls, envelope: MessageEnvelope) -> "GenerationThinking":
        p = envelope.payload
        return cls(
            chunk=p.get("chunk", ""),
            session_id=p.get("session_id", ""),
            query_id=p.get("query_id", ""),
        )

2.2 Publisher change

ZmqPublisher.publish() gains an overload that accepts a message object directly:

def publish(self, item: MessageEnvelope | BusMessage, *, sender_id: str = "", correlation_id: str = "") -> None:
    if isinstance(item, MessageEnvelope):
        envelope = item
    else:
        envelope = item.to_envelope(sender_id=sender_id, correlation_id=correlation_id)
    # ... existing send logic

Call sites become:

# Before
self._pub.publish(self._make_envelope(
    GENERATION_THINKING, "thinking",
    {"chunk": chunk.message.thinking, "session_id": session_id, "query_id": query_id},
    correlation_id, session_id,
))

# After
self._pub.publish(
    GenerationThinking(chunk=chunk.message.thinking, session_id=session_id, query_id=query_id),
    sender_id=self.id, correlation_id=correlation_id,
)

2.3 Consumer change

# Before
chunk = payload.get("chunk", "")
session_id = payload.get("session_id", "")

# After
msg = GenerationThinking.from_envelope(envelope)
chunk = msg.chunk
session_id = msg.session_id

2.4 Base classes

# Abstract base — all messages
class BusMessage:
    subject:      ClassVar[str]
    message_type: ClassVar[str]

    def to_envelope(self, sender_id: str, correlation_id: str) -> MessageEnvelope: ...

    @classmethod
    def from_envelope(cls, envelope: MessageEnvelope) -> "BusMessage": ...

# Tool call/result pairs share correlation_id + tool name
@dataclass
class ToolCall(BusMessage):
    tool:           str           # tool name — determines subject dynamically
    args:           dict
    correlation_id: str

    @property
    def subject(self) -> str:     # override: dynamic per tool name
        return f"tool.call.{self.tool}"

@dataclass
class ToolResult(BusMessage):
    tool:           str
    result:         str
    correlation_id: str

    @property
    def subject(self) -> str:
        return f"tool.result.{self.tool}"

For tools with fixed names, subclasses can lock the tool name as a class constant. ToolCall and ToolResult use a dynamic subject property because the subject includes the tool name.

2.5 message_type elimination

With message_type as a ClassVar on every message class, the separate MESSAGE_TYPE lookup dict discussed earlier is no longer needed. The class itself is the source of truth.

3. Message Catalogue

4. Impact Per Participant

5. Design Decisions to Confirm

5.1 Validation on from_envelope

Two options:

• Lenient (recommended): from_envelope uses .get() with defaults — same as today, but centralized. Never raises. Old envelopes from before the refactor still parse.
• Strict: raises ValueError on missing required fields. Catches wiring bugs earlier but breaks during incremental migration.

Recommended: lenient during migration, strict later once all sites are migrated.

5.2 ToolCall/ToolResult subject: class var vs property

Since the subject includes the tool name (tool.call.web_search), it can't be a fixed ClassVar. Options:

• Dynamic @property on the instance (shown above) — ClassVar[str] convention breaks
• Per-tool subclasses (WebSearchCall(ToolCall)) with fixed ClassVar — more classes but cleaner typing
• Keep as a single ToolCall class with a tool field and @property subject — simplest, matches how tool_dispatcher already works

Recommended: single ToolCall / ToolResult with @property subject.

5.3 _make_envelope in generator

Once all generator publish sites use message objects, _make_envelope can be removed. CompactionService currently receives it as a bound method parameter — that dependency is eliminated when CompactionService publishes a CompactionResult object directly via its own publisher.

5.4 Backward compatibility

The wire format (JSON payload) does not change — to_envelope calls asdict(self) which produces the same dict as today. Migration can be incremental, one participant at a time. If the wire format needs to change for any message, don't preserve backward compatibility — just update all sites together.

6. Migration Strategy

Because the wire format is unchanged, the refactor can be done incrementally:

1. Create src/local/protocol/messages.py with all dataclasses — no behavior changes yet
2. Update ZmqPublisher.publish() to accept BusMessage objects
3. Migrate publish sites one participant at a time (generator first — highest volume)
4. Migrate consume sites to use from_envelope()
5. Remove _make_envelope from GeneratorAgent once all its publish sites are migrated
6. Run full test suite + stories after each participant migration

7. Files to Create / Modify

• src/local/protocol/messages.py NEW
• src/local/transport/zmq_publisher.py MOD
• All participants listed in §4 above