test_codec.py

"""
ACCP Codec - Comprehensive Test Suite v2
Covers: typed values, envelope, canonical serialization, error frames,
        streaming, audit log, policy hooks, duplicate detection,
        domain profiles, adapter contract, session checkpointing.

Run: python poc/test_codec.py
 or: python -m pytest poc/test_codec.py -v
"""
import sys, os, time
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from poc.codec import (
    AccpEncoder, AccpDecoder, AgentMessage, Envelope, NULL,
    build_error_frame, build_stream_frame, build_ack_frame, _parse_frame,
)
from poc.registry import SchemaRegistry, ErrorCode
from poc.session import AccpSession, set_tokenizer_profile
from poc import profiles
from poc.adapter import AccpAdapter


def make_env():
    reg = SchemaRegistry()
    session = AccpSession()
    enc = AccpEncoder(reg)
    dec = AccpDecoder(reg)
    return enc, dec, session, reg


PASS = 0
FAIL = 0


def check(name: str, condition: bool, detail: str = "") -> None:
    global PASS, FAIL
    if condition:
        PASS += 1
        print(f"PASS  {name}")
    else:
        FAIL += 1
        print(f"FAIL  {name}" + (f" - {detail}" if detail else ""))


# ===========================================================================
# 1. Typed value round-trips
# ===========================================================================
def test_typed_values():
    from poc.codec import _serialize_value, _parse_typed_value
    cases = [
        (42,      "42"),
        (3.14,    "3.14"),
        (True,    "true"),
        (False,   "false"),
        (None,    "~"),
        (NULL,    "~"),
        ([1,2,3], "[1,2,3]"),
        ({"a":1, "b":2}, "{a:1,b:2}"),
        ("hello", "hello"),
        ("a:b",   r"a\:b"),
    ]
    for val, expected in cases:
        ser = _serialize_value(val)
        check(f"serialize {val!r}", ser == expected, f"got {ser!r}")
    # Parse back
    check("parse int",    _parse_typed_value("42") == 42)
    check("parse float",  abs(_parse_typed_value("3.14") - 3.14) < 1e-9)
    check("parse true",   _parse_typed_value("true") is True)
    check("parse false",  _parse_typed_value("false") is False)
    check("parse null",   repr(_parse_typed_value("~")) == "null")
    check("parse array",  _parse_typed_value("[1,2,3]") == [1, 2, 3])
    check("parse map",    _parse_typed_value("{a:1,b:2}") == {"a": 1, "b": 2})


# ===========================================================================
# 2. Canonical serialization - sorted field order
# ===========================================================================
def test_canonical_order():
    enc, _, session, _ = make_env()
    msg = AgentMessage("agent", "req", "test",
                       payload={"z": 1, "a": 2, "m": 3})
    frame = enc.encode(msg, session)
    idx_a = frame.raw.index("a:2")
    idx_m = frame.raw.index("m:3")
    idx_z = frame.raw.index("z:1")
    check("canonical sort a<m<z", idx_a < idx_m < idx_z)


# ===========================================================================
# 3. Envelope fields present and decoded
# ===========================================================================
def test_envelope_roundtrip():
    enc, _, enc_session, _ = make_env()
    msg = AgentMessage("planner", "req", "schedule",
                       payload={"task": "auth"},
                       envelope=Envelope(correlation_id="corr123", ttl=60))
    frame = enc.encode(msg, enc_session)
    check("envelope in frame", "mid:" in frame.raw)
    check("seq in frame", "seq:" in frame.raw)
    check("cid in frame", "cid:corr123" in frame.raw)
    # Decode on a fresh session (no shared seen-msg-ids)
    dec_session = AccpSession()
    dec = AccpDecoder(SchemaRegistry())
    restored = dec.decode(frame.raw, dec_session)
    check("envelope msg_id decoded", len(restored.envelope.msg_id) == 12)
    check("envelope corr decoded",   restored.envelope.correlation_id == "corr123")
    check("envelope seq decoded",    restored.envelope.sequence == 1)


# ===========================================================================
# 4. Duplicate detection
# ===========================================================================
def test_duplicate_detection():
    enc, dec, session, _ = make_env()
    msg = AgentMessage("a", "req", "op", payload={"x": 1})
    frame = enc.encode(msg, session)
    try:
        dec.decode(frame.raw, session)
        check("duplicate raises", False, "no exception raised")
    except ValueError as e:
        check("duplicate raises", "Duplicate" in str(e))


# ===========================================================================
# 5. Policy hooks
# ===========================================================================
def test_policy_hook():
    enc, _, enc_session, reg = make_env()
    dec = AccpDecoder(reg)
    def deny_fail_intent(msg: AgentMessage) -> AgentMessage:
        if msg.intent == "fail":
            raise PermissionError("Policy denied: fail intent not allowed here")
        return msg
    dec.add_policy_hook(deny_fail_intent)
    fail_msg = AgentMessage("agent", "fail", "op", payload={"err": "boom"})
    frame = enc.encode(fail_msg, enc_session)
    # Decode on a fresh session so dup detection doesn't fire first
    dec_session = AccpSession()
    try:
        dec.decode(frame.raw, dec_session)
        check("policy hook fires", False, "no exception raised")
    except PermissionError:
        check("policy hook fires", True)


# ===========================================================================
# 6. Error frame builder
# ===========================================================================
def test_error_frame():
    enc, dec, session, _ = make_env()
    frame = build_error_frame("agent", ErrorCode.TOOL_NOT_FOUND,
                              "tool missing", session=session, encoder=enc)
    check("error frame intent=fail", frame.intent == "fail")
    check("error frame op=error",    frame.operation == "error")
    check("error_code in raw",       "E4001" in frame.raw or "code" in frame.raw)


# ===========================================================================
# 7. Streaming frames
# ===========================================================================
def test_streaming():
    enc, dec, session, _ = make_env()
    chunks = ["Hello", "world", "!"]
    cid = "stream_abc"
    frames = []
    for i, chunk in enumerate(chunks):
        is_final = (i == len(chunks) - 1)
        f = build_stream_frame("streamer", "infer", i, len(chunks),
                               chunk, is_final, cid, session, enc)
        frames.append(f)
    check("3 stream frames built", len(frames) == 3)
    check("all intent=stream", all(f.intent == "stream" for f in frames))
    # decode last chunk independently (new session to avoid dup detection)
    s2 = AccpSession()
    dec2 = AccpDecoder(SchemaRegistry())
    last = dec2.decode(frames[-1].raw, s2)
    check("is_final restored", last.payload.get("is_final") is True)


# ===========================================================================
# 8. ACK frame builder
# ===========================================================================
def test_ack_frame():
    enc, dec, session, _ = make_env()
    frame = build_ack_frame("agent", "abc123msg", session, enc)
    check("ack intent", frame.intent == "ack")
    check("ack payload", frame.payload.get("acked") == "abc123msg")


# ===========================================================================
# 9. Audit log populated
# ===========================================================================
def test_audit_log():
    enc, dec, session, _ = make_env()
    msg = AgentMessage("a1", "req", "op", payload={"x": 1})
    frame = enc.encode(msg, session)
    s2 = AccpSession()
    dec2 = AccpDecoder(SchemaRegistry())
    dec2.decode(frame.raw, s2)
    check("encode audit entry",  len(session.audit_log) == 1)
    check("decode audit entry",  len(s2.audit_log) == 1)
    check("audit direction enc", session.audit_log[0].direction == "encode")
    check("audit direction dec", s2.audit_log[0].direction == "decode")


# ===========================================================================
# 10. Tokenizer profiles
# ===========================================================================
def test_tokenizer_profiles():
    from poc.session import _estimate_tokens, set_tokenizer_profile
    text = "A" * 350
    set_tokenizer_profile("gpt")
    t_gpt = _estimate_tokens(text)
    set_tokenizer_profile("claude")
    t_claude = _estimate_tokens(text)
    set_tokenizer_profile("gemini")
    t_gemini = _estimate_tokens(text)
    # All should be close but not identical
    check("tokenizer profiles differ", not (t_gpt == t_claude == t_gemini),
          f"gpt={t_gpt} claude={t_claude} gemini={t_gemini}")
    set_tokenizer_profile("claude")  # restore default


# ===========================================================================
# 11. Warm budget enforcement
# ===========================================================================
def test_warm_budget():
    from poc.session import WarmStateStore, WARM_TOKEN_MAX_PER_CHECKPOINT
    store = WarmStateStore()
    huge = {"k" + str(i): "v" * 30 for i in range(100)}
    try:
        store.put("ckpt_big", huge)
        check("warm budget raises on oversize", False, "no exception")
    except ValueError:
        check("warm budget raises on oversize", True)


# ===========================================================================
# 12. TTL expiry
# ===========================================================================
def test_ttl_expiry():
    import time as _time
    from poc.session import InMemoryColdStore
    store = InMemoryColdStore()
    # Use a negative TTL to guarantee expiry before the get() call
    store.put("k1", "val", ttl_seconds=-1)
    check("cold TTL expiry", store.get("k1") is None)


# ===========================================================================
# 13. Hot-state accumulation
# ===========================================================================
def test_hot_state_accumulation():
    enc, _, session, _ = make_env()
    for i in range(3):
        msg = AgentMessage("planner", "req", "schedule",
                           payload={"task": f"task_{i}", "priority": "high"})
        enc.encode(msg, session)
    slot = "planner.schedule"
    check("hot state accumulated", slot in session.hot_state)
    check("hot state latest task", session.hot_state[slot].get("task") == "task_2")


# ===========================================================================
# 14. Domain profiles
# ===========================================================================
def test_profiles_chat():
    enc, dec, session, _ = make_env()
    msg = profiles.chat("assistant", "Here are the results", turn=2)
    frame = enc.encode(msg, session)
    s2 = AccpSession()
    restored = AccpDecoder(SchemaRegistry()).decode(frame.raw, s2)
    check("chat content", restored.payload.get("content") == "Here are the results")
    check("chat turn",    restored.payload.get("turn") == 2)
    check("chat role default restored", restored.payload.get("role") == "assistant")


def test_profiles_tool_call():
    enc, _, session, _ = make_env()
    msg = profiles.tool_call("orch", "web_search",
                             {"query": "ACCP", "max_results": 5}, "cid1")
    frame = enc.encode(msg, session)
    check("tool_call intent=req", "req" in frame.raw)
    check("tool_name in frame",   "web_search" in frame.raw)


def test_profiles_transaction():
    enc, _, session, _ = make_env()
    msg = profiles.transaction("payments", "txn_001", 142.50, "acct_9876")
    frame = enc.encode(msg, session)
    check("transaction frame", "transaction" in frame.raw or "txn" in frame.raw)
    check("amount serialized", "142.5" in frame.raw)
    check("USD omitted (default)", "USD" not in frame.raw)


def test_profiles_workflow():
    enc, dec, session, _ = make_env()
    msg = profiles.workflow("planner", "impl_auth", "@dev",
                            priority="high", deadline="sprint_14")
    frame = enc.encode(msg, session)
    s2 = AccpSession()
    restored = AccpDecoder(SchemaRegistry()).decode(frame.raw, s2)
    check("workflow priority", restored.payload.get("priority") == "high")
    check("workflow deps default", restored.payload.get("deps") == [])


# ===========================================================================
# 15. Adapter base class
# ===========================================================================
def test_adapter():
    sent = []
    class TestAdapter(AccpAdapter):
        def on_message_out(self, frame):
            sent.append(frame.raw)
        def on_message_in(self, msg, raw):
            pass

    adapter = TestAdapter("test_agent")
    sync_frame = adapter.on_session_start()
    check("adapter sync frame", "sync" in sync_frame.raw)
    # on_session_start encodes but does NOT call on_message_out; send() does
    msg = profiles.chat("test_agent", "Hello from adapter", turn=1)
    adapter.send(msg)
    check("adapter send captured", len(sent) == 1)  # only chat send goes via on_message_out

    err_frame = adapter.on_error(ErrorCode.TIMEOUT, "connection timed out")
    check("adapter error frame", err_frame.intent == "fail")


# ===========================================================================
# 16. Registry - schema profiles available
# ===========================================================================
def test_registry_profiles():
    reg = SchemaRegistry()
    tool_schemas = reg.schemas_for_profile("tool")
    txn_schemas  = reg.schemas_for_profile("transaction")
    chat_schemas = reg.schemas_for_profile("chat")
    check("tool profile schemas",        len(tool_schemas) >= 1)
    check("transaction profile schemas", len(txn_schemas) >= 1)
    check("chat profile schemas",        len(chat_schemas) >= 1)


# ===========================================================================
# 17. Original spec frame examples still parse correctly
# ===========================================================================
def test_spec_examples():
    frames = [
        "@research>done:analyze{d:q3_sales|f:[rev:-12%QoQ,churn:+3.2%]|nx:@strategy:plan}",
        "@planner>req:schedule{pri:high|task:impl_auth_module|when:sprint_14|who:@dev_team}",
        "@analyst>qry:lookup{fmt:summary|q:revenue_by_region|src:$ctx.sales_db}",
        "@orchestrator>sync:state{delta:{task_3:done,task_4:wip}|v:7}",
        "@data_agent>fail:fetch{err:timeout_30s|retry:3|src:api.crm}",
    ]
    for raw in frames:
        try:
            f = _parse_frame(raw)
            check(f"parse spec frame {f.agent_id}>{f.intent}", True)
        except Exception as e:
            check(f"parse spec frame", False, str(e))


# ===========================================================================
# Runner
# ===========================================================================
if __name__ == "__main__":
    suites = [
        test_typed_values, test_canonical_order, test_envelope_roundtrip,
        test_duplicate_detection, test_policy_hook, test_error_frame,
        test_streaming, test_ack_frame, test_audit_log,
        test_tokenizer_profiles, test_warm_budget, test_ttl_expiry,
        test_hot_state_accumulation,
        test_profiles_chat, test_profiles_tool_call,
        test_profiles_transaction, test_profiles_workflow,
        test_adapter, test_registry_profiles, test_spec_examples,
    ]
    for suite in suites:
        try:
            suite()
        except Exception as e:
            print(f"ERROR in {suite.__name__}: {e}")
            FAIL += 1
    total = PASS + FAIL
    print(f"\n{PASS}/{total} passed  ({FAIL} failed)")
    sys.exit(FAIL)