Release v0.1.5

tests/test_runtime.py

@@ -7,6 +7,8 @@ from micro_agents import (
     LLM,
     LlmMessage,
     Media,
+    MessageAdded,
+    SessionState,
     Tool,
     ToolCallFailed,
     ToolCallMessage,
@@ -72,6 +74,39 @@ class SessionRuntimeTests(unittest.IsolatedAsyncioTestCase):
         responses = [m for m in session.conversation.messages if isinstance(m, ToolCallResponseMessage)]
         self.assertEqual('{"echo":"hello"}', responses[0].text())
 
+    async def test_user_prompt_receives_context_with_source_text(self):
+        llm = ScriptedLLM([[LlmMessage("done")]])
+        context = {"user_id": 42, "source_text": "stale"}
+        prompt_contexts = []
+
+        async def user_prompt(prompt_context):
+            prompt_contexts.append(prompt_context)
+            return f"<user id={prompt_context['user_id']}>{prompt_context['source_text']}</user>"
+
+        session = Aggregator(llm, user_prompt=user_prompt).create_session(context=context)
+
+        result = await session.run("original")
+
+        self.assertTrue(result.succeeded)
+        user_message = session.conversation.messages[0]
+        self.assertIsInstance(user_message, UserMessage)
+        assert isinstance(user_message, UserMessage)
+        self.assertEqual("<user id=42>original</user>", user_message.message())
+        self.assertEqual([{"user_id": 42, "source_text": "original"}], prompt_contexts)
+        self.assertEqual({"user_id": 42, "source_text": "stale"}, context)
+
+    async def test_user_prompt_keeps_media_on_the_formatted_message(self):
+        llm = ScriptedLLM([[LlmMessage("seen")]])
+        media = Media("image", data=b"image", mime_type="image/png", filename="image.png")
+        session = Aggregator(llm, user_prompt=lambda context: f"Input: {context['source_text']}").create_session()
+
+        await session.run("describe", media=[media])
+
+        user_message = session.conversation.messages[0]
+        assert isinstance(user_message, UserMessage)
+        self.assertEqual("Input: describe", user_message.message())
+        self.assertEqual([media], user_message.media())
+
     async def test_permission_is_checked_again_before_execution(self):
         def restricted():
             raise AssertionError("A denied tool must not execute")
@@ -171,6 +206,71 @@ class SessionRuntimeTests(unittest.IsolatedAsyncioTestCase):
         assert isinstance(user_message, UserMessage)
         self.assertEqual([media], user_message.media())
 
+    async def test_session_state_can_be_saved_and_restored(self):
+        media = Media("image", data=b"image", mime_type="image/png", filename="image.png")
+        context = {"user_id": 42}
+        first_llm = ScriptedLLM([[LlmMessage("first answer")]])
+        first_session = Aggregator(first_llm).create_session(context=context, session_name="chat-42")
+        await first_session.run("first question", media=[media])
+
+        state = first_session.save()
+        loaded_state = SessionState.from_json(state.to_json())
+        restored_llm = ScriptedLLM([[LlmMessage("second answer")]])
+        restored_session = Aggregator(restored_llm).create_session(state=loaded_state)
+        context["user_id"] = 7
+
+        events = [event async for event in restored_session.stream("second question")]
+
+        self.assertEqual("chat-42", restored_session.id)
+        self.assertEqual("chat-42", restored_session.name)
+        self.assertEqual({"user_id": 42}, restored_session.context)
+        self.assertEqual(4, len(restored_session.conversation.messages))
+        restored_user_message = restored_session.conversation.messages[0]
+        assert isinstance(restored_user_message, UserMessage)
+        self.assertEqual([media], restored_user_message.media())
+        added_messages = [event.message for event in events if isinstance(event, MessageAdded)]
+        added_user_message = added_messages[0]
+        assert isinstance(added_user_message, UserMessage)
+        self.assertEqual("second question", added_user_message.message())
+
+    def test_session_state_preserves_tool_and_reasoning_messages(self):
+        messages = [
+            ReasoningMessage("thinking", provider_item={"type": "reasoning", "id": "reasoning-1"}),
+            ToolCallMessage("call-1", "lookup", {"query": "value"}),
+            ToolCallResponseMessage("call-1", "result", is_error=True, error_code="lookup_failed"),
+        ]
+
+        restored = SessionState.from_json(SessionState.capture("id", None, {}, messages).to_json()).restored_messages()
+
+        reasoning, call, response = restored
+        assert isinstance(reasoning, ReasoningMessage)
+        assert isinstance(call, ToolCallMessage)
+        assert isinstance(response, ToolCallResponseMessage)
+        self.assertEqual({"type": "reasoning", "id": "reasoning-1"}, reasoning.provider_item())
+        self.assertEqual({"query": "value"}, call.arguments())
+        self.assertTrue(response.is_error())
+        self.assertEqual("lookup_failed", response.error_code())
+
+    def test_create_session_accepts_state_dictionary(self):
+        state = SessionState.capture(
+            "saved-id",
+            None,
+            {"tenant": "example"},
+            [UserMessage("hello"), LlmMessage("hi")],
+        )
+
+        session = Aggregator(ScriptedLLM([])).create_session(state=state.to_dict())
+
+        self.assertEqual("saved-id", session.id)
+        self.assertEqual({"tenant": "example"}, session.context)
+        self.assertEqual(2, len(session.conversation.messages))
+
+    def test_state_cannot_be_mixed_with_initial_session_data(self):
+        state = SessionState.capture("saved-id", None, {}, [])
+
+        with self.assertRaises(ValueError):
+            Aggregator(ScriptedLLM([])).create_session(state=state, context={})
+
     def test_openai_serializes_inline_image(self):
         media = Media("image", data=b"image", mime_type="image/png", filename="image.png")