AgentCoreMemorySessionManager: BidiAgent transcripts are never stored (MessageUpdatedEvent not handled) #716

"""BidiAgent + AgentCoreMemorySessionManager: what gets stored, and does it hold up the audio? No AWS calls."""
import asyncio, base64, time

from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager
from strands.bidi.agent import BidiAgent
from strands.bidi.models.model import BidiModel
from strands.bidi.types import events as ev

LATENCY = 0.1  # fake memory data plane latency per call
calls = []


class FakeDataPlane:
    class meta:
        region_name = "us-east-1"

    def list_events(self, **kw):
        time.sleep(LATENCY)
        return {"events": []}

    def create_event(self, **kw):
        time.sleep(LATENCY)
        for p in kw["payload"]:
            calls.append(("conversational", p["conversational"]["role"], p["conversational"]["content"]["text"][:60])
                         if "conversational" in p else ("blob", kw.get("metadata", {}).get("stateType", {}).get("stringValue")))
        return {"event": {"eventId": str(len(calls))}}


class FakeSession:
    region_name = "us-east-1"

    def client(self, *args, **kwargs):
        return FakeDataPlane()


class FakeNovaLikeModel(BidiModel):
    """One user turn, then one assistant turn with audio, in the event order Nova Sonic produces."""

    def __init__(self):
        self.config, self.marks, self.stopped = {"model_id": "fake"}, {}, asyncio.Event()

    def get_config(self):
        return self.config

    def update_config(self, **kw):
        self.config.update(kw)

    async def start(self, system_prompt=None, tools=None, messages=None, **kw):
        pass

    async def stop(self):
        self.stopped.set()

    async def send(self, content):
        pass

    async def receive(self):
        yield ev.BidiResponseStartEvent("r1")
        yield ev.BidiTranscriptStartEvent("user", "u1")
        yield ev.BidiTranscriptDeltaEvent("I need a hotel in London.", "user", "u1")
        yield ev.BidiTranscriptStopEvent("user", "u1")
        self.marks["assistant_start"] = time.monotonic()
        yield ev.BidiTranscriptStartEvent("assistant", "a1")
        yield ev.BidiTranscriptDeltaEvent("Lovely, when are you travelling?", "assistant", "a1")
        yield ev.BidiAudioStartEvent("au1")
        yield ev.BidiAudioDeltaEvent(base64.b64encode(b"\0\0" * 480).decode(), "pcm", 24000, 1, "au1")
        yield ev.BidiAudioStopEvent("au1")
        yield ev.BidiTranscriptStopEvent("assistant", "a1")
        yield ev.BidiResponseStopEvent("r1")
        await self.stopped.wait()


async def main():
    config = AgentCoreMemoryConfig(memory_id="mem-1", session_id="session-1", actor_id="user-1", async_mode=True)
    manager = AgentCoreMemorySessionManager(config, region_name="us-east-1", boto_session=FakeSession())
    model = FakeNovaLikeModel()
    agent = BidiAgent(model=model, system_prompt="x", session_manager=manager,
                      messages=[{"role": "assistant", "content": [{"text": "Hi, how can I help?"}]}])
    await agent.start()
    first_audio = None
    async for event in agent.receive():
        if event.get("type") == "bidi_audio_delta" and first_audio is None:
            first_audio = time.monotonic()
        if event.get("type") == "bidi_response_stop":
            break
    await agent.stop()
    print("history:", [(m["role"], m["content"][0]["text"]) for m in agent.messages])
    print("stored in AgentCore Memory:")
    for c in calls:
        print("  ", c)
    print(f"assistant transcript start -> first audio event: {(first_audio - model.marks['assistant_start']) * 1000:.0f} ms "
          f"(fake write latency {LATENCY * 1000:.0f} ms)")


asyncio.run(main())