Release v0.1.5

README.md

@@ -20,11 +20,16 @@ def master_prompt(context: dict[str, object]) -> str:
     return f"You are a concise support bot. Reply in {context.get('language', 'en')}."
 
 
+def user_prompt(context: dict[str, object]) -> str:
+    return f"<user_message>\n{context['source_text']}\n</user_message>"
+
+
 llm = OpenAIBasedLLM(api_url="https://api.openai.com/v1", api_key="...", model="gpt-5")
 agent = Aggregator(
     llm,
     BotTools([Tool(current_weather)]),
     master_prompt=master_prompt,
+    user_prompt=user_prompt,
 )
 session = agent.create_session(context={"user_id": 42, "role": "member", "language": "ru"})
 
@@ -56,6 +61,24 @@ result = await session.run(
 
 Do not interpolate secrets, API keys, or internal-only values into `master_prompt`: it receives the full live context dict.
 
+`user_prompt` can synchronously or asynchronously normalize, wrap, or enrich every incoming user message before it is added to the conversation. It receives a copy of the session context with `source_text` set to the original text passed to `run()` or `stream()`; this reserved value overrides any `source_text` already present in the context without mutating the context itself. The returned string becomes the user message, while any `media` attachments stay attached to it.
+
+## Saving and restoring sessions
+
+Save the application context and complete conversation, then create a new session from that state:
+
+```python
+from micro_agents import SessionState
+
+state = session.save()
+await database.save(state.session_id, state.to_json())
+
+loaded = SessionState.from_json(await database.load(state.session_id))
+session = agent.create_session(state=loaded)
+```
+
+`create_session()` also accepts the result of `state.to_dict()` directly. Restoring preserves the session ID, optional name, messages, inline media, tool calls, tool responses, and provider reasoning items. The restored context and messages are fresh containers, so continuing the restored session does not mutate the saved state. Context values must be JSON-compatible when `to_json()` is used; pass a `default` converter for values such as `Path`, for example `state.to_json(default=str)`.
+
 ## Dynamic tools
 
 Use `DynamicToolManager` when integrations can be connected or disconnected while the application is running:

packages/micro-agents-tools/pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "micro-agents-tools"
-version = "0.1.4"
+version = "0.1.5"
 description = "Basic opt-in tools for micro-agents"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "micro-agents>=0.1.4",
+    "micro-agents>=0.1.5",
 ]
 
 [project.urls]

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "micro-agents"
-version = "0.1.4"
+version = "0.1.5"
 description = "Small async building blocks for embedding LLM tool loops in Python applications"
 readme = "README.md"
 authors = [

src/micro_agents/__init__.py

@@ -12,10 +12,12 @@ from micro_agents.aggregator import (
     ToolCallFailed,
     ToolCallFinished,
     ToolCallStarted,
+    UserPrompt,
 )
 from micro_agents.context import Injected, SessionContext
 from micro_agents.conversation import Conversation, LlmMessage, Media, ToolCallMessage, ToolCallResponseMessage, UserMessage
 from micro_agents.llm import LLM, OpenAIBasedLLM
+from micro_agents.state import SessionState
 from micro_agents.tools import DynamicToolManager, Tool, ToolManager, ToolOutput
 
 __all__ = [
@@ -37,6 +39,7 @@ __all__ = [
     "SessionContext",
     "SessionEvent",
     "SessionResult",
+    "SessionState",
     "Tool",
     "ToolCallFailed",
     "ToolCallFinished",
@@ -46,6 +49,7 @@ __all__ = [
     "ToolManager",
     "ToolOutput",
     "UserMessage",
+    "UserPrompt",
 ]
 
 

src/micro_agents/aggregator.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import inspect
-from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping
 from dataclasses import dataclass
 from typing import Protocol
 from uuid import uuid4
@@ -10,11 +10,13 @@ from uuid import uuid4
 from micro_agents.context import SessionContext
 from micro_agents.conversation import Conversation, LlmMessage, Media, Message, ToolCallMessage, ToolCallResponseMessage, UserMessage
 from micro_agents.llm import LLM
+from micro_agents.state import SessionState
 from micro_agents.tools import ToolManager
 
 
 Context = SessionContext
 MasterPrompt = Callable[[SessionContext], str | Awaitable[str]]
+UserPrompt = Callable[[SessionContext], str | Awaitable[str]]
 
 
 class EventSubscriber(Protocol):
@@ -79,8 +81,10 @@ class Session:
         *,
         context: Context | None = None,
         master_prompt: MasterPrompt | None = None,
+        user_prompt: UserPrompt | None = None,
         event_subscriber: EventSubscriber | object | None = None,
         session_name: str | None = None,
+        session_id: str | None = None,
         messages: Iterable[Message] | None = None,
         max_tool_rounds: int = 16,
         thinking: bool | None = None,
@@ -93,9 +97,10 @@ class Session:
         # quotas, or locale between turns.
         self.context: Context = context if context is not None else {}
         self.name = session_name
-        self.id = session_name or str(uuid4())
+        self.id = session_id or session_name or str(uuid4())
         self.conversation = Conversation(list(messages or ()))
         self._master_prompt = master_prompt
+        self._user_prompt = user_prompt
         self._event_subscriber = event_subscriber
         self._max_tool_rounds = max_tool_rounds
         self._thinking = thinking
@@ -117,12 +122,17 @@ class Session:
                 return SessionResult(None, self.conversation, event.error)
         return SessionResult(None, self.conversation, RuntimeError("Session ended without a result"))
 
-    async def _run_events(self, user_text: str, media: Iterable[Media] | None) -> AsyncIterator[SessionEvent]:
-        user_message = UserMessage(user_text, list(media or ()))
-        self.conversation.add_message(user_message)
-        yield MessageAdded(self, user_message)
+    def save(self) -> SessionState:
+        """Capture the context and transcript needed to resume this session."""
+        return SessionState.capture(self.id, self.name, self.context, self.conversation.messages)
 
+    async def _run_events(self, user_text: str, media: Iterable[Media] | None) -> AsyncIterator[SessionEvent]:
         try:
+            formatted_user_text = await self._build_user_prompt(user_text)
+            user_message = UserMessage(formatted_user_text, list(media or ()))
+            self.conversation.add_message(user_message)
+            yield MessageAdded(self, user_message)
+
             for _ in range(self._max_tool_rounds):
                 instructions = await self._build_master_prompt()
                 tools = await self._tool_manager.available_tools(self.context)
@@ -174,6 +184,18 @@ class Session:
             raise TypeError("master_prompt must return str or awaitable[str]")
         return prompt
 
+    async def _build_user_prompt(self, source_text: str) -> str:
+        if self._user_prompt is None:
+            return source_text
+        prompt_context = dict(self.context)
+        prompt_context["source_text"] = source_text
+        prompt = self._user_prompt(prompt_context)
+        if not isinstance(prompt, str):
+            prompt = await prompt
+        if not isinstance(prompt, str):
+            raise TypeError("user_prompt must return str or awaitable[str]")
+        return prompt
+
     async def _notify(self, event: SessionEvent) -> None:
         if self._event_subscriber is None:
             return
@@ -201,12 +223,14 @@ class Aggregator:
         tool_manager: ToolManager | None = None,
         *,
         master_prompt: MasterPrompt | None = None,
+        user_prompt: UserPrompt | None = None,
         max_tool_rounds: int = 16,
         thinking: bool | None = None,
     ):
         self._llm = llm
         self._tool_manager = tool_manager or ToolManager()
         self._master_prompt = master_prompt
+        self._user_prompt = user_prompt
         self._max_tool_rounds = max_tool_rounds
         self._thinking = thinking
 
@@ -217,16 +241,28 @@ class Aggregator:
         event_subscriber: EventSubscriber | object | None = None,
         session_name: str | None = None,
         messages: Iterable[Message] | None = None,
+        state: SessionState | Mapping[str, object] | None = None,
         llm: LLM | None = None,
         thinking: bool | None = None,
     ) -> Session:
+        session_id: str | None = None
+        if state is not None:
+            if context is not None or session_name is not None or messages is not None:
+                raise ValueError("context, session_name, and messages cannot be combined with state")
+            loaded_state = state if isinstance(state, SessionState) else SessionState.from_dict(state)
+            context = dict(loaded_state.context)
+            session_name = loaded_state.session_name
+            session_id = loaded_state.session_id
+            messages = loaded_state.restored_messages()
         return Session(
             self._llm if llm is None else llm,
             self._tool_manager,
             context=context,
             master_prompt=self._master_prompt,
+            user_prompt=self._user_prompt,
             event_subscriber=event_subscriber,
             session_name=session_name,
+            session_id=session_id,
             messages=messages,
             max_tool_rounds=self._max_tool_rounds,
             thinking=self._thinking if thinking is None else thinking,

src/micro_agents/state.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import base64
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass
+from typing import cast
+
+import orjson
+
+from micro_agents.context import SessionContext
+from micro_agents.conversation import (
+    LlmMessage,
+    Media,
+    MediaDetail,
+    MediaType,
+    Message,
+    ReasoningMessage,
+    ToolCallMessage,
+    ToolCallResponseMessage,
+    UserMessage,
+)
+
+
+_STATE_VERSION = 1
+
+
+@dataclass(frozen=True, slots=True)
+class SessionState:
+    """Serializable snapshot used to resume a session later."""
+
+    session_id: str
+    session_name: str | None
+    context: dict[str, object]
+    messages: tuple[dict[str, object], ...]
+
+    @classmethod
+    def capture(
+        cls,
+        session_id: str,
+        session_name: str | None,
+        context: SessionContext,
+        messages: Iterable[Message],
+    ) -> "SessionState":
+        return cls(
+            session_id=session_id,
+            session_name=session_name,
+            context=dict(context),
+            messages=tuple(_serialize_message(message) for message in messages),
+        )
+
+    def restored_messages(self) -> list[Message]:
+        """Create fresh message objects for a restored conversation."""
+        return [_deserialize_message(message) for message in self.messages]
+
+    def to_dict(self) -> dict[str, object]:
+        """Return storage-friendly state; context values remain application-owned."""
+        return {
+            "version": _STATE_VERSION,
+            "session_id": self.session_id,
+            "session_name": self.session_name,
+            "context": dict(self.context),
+            "messages": [dict(message) for message in self.messages],
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, object]) -> "SessionState":
+        version = data.get("version")
+        if version != _STATE_VERSION:
+            raise ValueError(f"Unsupported session state version: {version!r}")
+        session_id = _required_str(data, "session_id")
+        session_name_value = data.get("session_name")
+        if session_name_value is not None and not isinstance(session_name_value, str):
+            raise TypeError("session_name must be str or None")
+        context_value = data.get("context")
+        if not isinstance(context_value, Mapping):
+            raise TypeError("context must be a mapping")
+        messages_value = data.get("messages")
+        if not isinstance(messages_value, (list, tuple)):
+            raise TypeError("messages must be a list")
+        serialized_messages: list[dict[str, object]] = []
+        for message in messages_value:
+            if not isinstance(message, Mapping):
+                raise TypeError("Every saved message must be a mapping")
+            normalized = dict(cast(Mapping[str, object], message))
+            _deserialize_message(normalized)
+            serialized_messages.append(normalized)
+        return cls(
+            session_id=session_id,
+            session_name=session_name_value,
+            context=dict(cast(Mapping[str, object], context_value)),
+            messages=tuple(serialized_messages),
+        )
+
+    def to_json(self, *, default: Callable[[object], object] | None = None) -> bytes:
+        """Encode state as JSON, optionally converting custom context values."""
+        return orjson.dumps(self.to_dict(), default=default)
+
+    @classmethod
+    def from_json(cls, data: str | bytes | bytearray | memoryview) -> "SessionState":
+        decoded = orjson.loads(data)
+        if not isinstance(decoded, dict):
+            raise TypeError("Session state JSON must contain an object")
+        return cls.from_dict(cast(dict[str, object], decoded))
+
+
+def _serialize_message(message: Message) -> dict[str, object]:
+    if isinstance(message, UserMessage):
+        return {"type": "user", "text": message.message(), "media": [_serialize_media(m) for m in message.media()]}
+    if isinstance(message, ReasoningMessage):
+        return {"type": "reasoning", "text": message.text(), "provider_item": message.provider_item()}
+    if isinstance(message, LlmMessage):
+        return {"type": "llm", "text": message.message()}
+    if isinstance(message, ToolCallMessage):
+        return {
+            "type": "tool_call",
+            "unique_id": message.unique_id(),
+            "tool_name": message.tool_name(),
+            "arguments": dict(message.arguments()),
+        }
+    if isinstance(message, ToolCallResponseMessage):
+        return {
+            "type": "tool_response",
+            "unique_id": message.unique_id(),
+            "text": message.text(),
+            "attachments": [_serialize_media(media) for media in message.attachments()],
+            "is_error": message.is_error(),
+            "error_code": message.error_code(),
+        }
+    raise TypeError(f"Cannot save message type {type(message).__name__}")
+
+
+def _deserialize_message(data: Mapping[str, object]) -> Message:
+    message_type = _required_str(data, "type")
+    if message_type == "user":
+        return UserMessage(_required_str(data, "text"), _deserialize_media_list(data.get("media"), "media"))
+    if message_type == "reasoning":
+        provider_item = data.get("provider_item")
+        if provider_item is not None and not isinstance(provider_item, Mapping):
+            raise TypeError("provider_item must be a mapping or None")
+        return ReasoningMessage(
+            _required_str(data, "text"), provider_item=cast(Mapping[str, object] | None, provider_item)
+        )
+    if message_type == "llm":
+        return LlmMessage(_required_str(data, "text"))
+    if message_type == "tool_call":
+        arguments = data.get("arguments")
+        if not isinstance(arguments, Mapping):
+            raise TypeError("arguments must be a mapping")
+        return ToolCallMessage(
+            _required_str(data, "unique_id"),
+            _required_str(data, "tool_name"),
+            dict(cast(Mapping[str, object], arguments)),
+        )
+    if message_type == "tool_response":
+        is_error = data.get("is_error", False)
+        if not isinstance(is_error, bool):
+            raise TypeError("is_error must be bool")
+        error_code = data.get("error_code")
+        if error_code is not None and not isinstance(error_code, str):
+            raise TypeError("error_code must be str or None")
+        return ToolCallResponseMessage(
+            _required_str(data, "unique_id"),
+            _required_str(data, "text"),
+            attachments=_deserialize_media_list(data.get("attachments"), "attachments"),
+            is_error=is_error,
+            error_code=error_code,
+        )
+    raise ValueError(f"Unknown saved message type: {message_type!r}")
+
+
+def _serialize_media(media: Media) -> dict[str, object]:
+    return {
+        "kind": media.kind,
+        "data": base64.b64encode(media.data).decode("ascii") if media.data is not None else None,
+        "url": media.url,
+        "file_id": media.file_id,
+        "filename": media.filename,
+        "mime_type": media.mime_type,
+        "detail": media.detail,
+    }
+
+
+def _deserialize_media(data: Mapping[str, object]) -> Media:
+    kind = _required_str(data, "kind")
+    if kind not in ("image", "file", "audio"):
+        raise ValueError(f"Unknown media kind: {kind!r}")
+    detail = data.get("detail", "auto")
+    if detail not in ("auto", "low", "high"):
+        raise ValueError(f"Unknown media detail: {detail!r}")
+    encoded_data = data.get("data")
+    if encoded_data is not None and not isinstance(encoded_data, str):
+        raise TypeError("Media data must be a base64 string or None")
+    binary_data = base64.b64decode(encoded_data, validate=True) if encoded_data is not None else None
+    return Media(
+        cast(MediaType, kind),
+        data=binary_data,
+        url=_optional_str(data, "url"),
+        file_id=_optional_str(data, "file_id"),
+        filename=_optional_str(data, "filename"),
+        mime_type=_optional_str(data, "mime_type"),
+        detail=cast(MediaDetail, detail),
+    )
+
+
+def _deserialize_media_list(value: object, field_name: str) -> list[Media]:
+    if value is None:
+        return []
+    if not isinstance(value, (list, tuple)):
+        raise TypeError(f"{field_name} must be a list")
+    result: list[Media] = []
+    for item in value:
+        if not isinstance(item, Mapping):
+            raise TypeError(f"Every {field_name} item must be a mapping")
+        result.append(_deserialize_media(cast(Mapping[str, object], item)))
+    return result
+
+
+def _required_str(data: Mapping[str, object], key: str) -> str:
+    value = data.get(key)
+    if not isinstance(value, str):
+        raise TypeError(f"{key} must be str")
+    return value
+
+
+def _optional_str(data: Mapping[str, object], key: str) -> str | None:
+    value = data.get(key)
+    if value is not None and not isinstance(value, str):
+        raise TypeError(f"{key} must be str or None")
+    return value

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")
 

uv.lock

@@ -157,7 +157,7 @@ wheels = [
 
 [[package]]
 name = "micro-agents"
-version = "0.1.4"
+version = "0.1.5"
 source = { editable = "." }
 dependencies = [
     { name = "openai" },
@@ -180,7 +180,7 @@ dev = [{ name = "micro-agents-tools", editable = "packages/micro-agents-tools" }
 
 [[package]]
 name = "micro-agents-tools"
-version = "0.1.4"
+version = "0.1.5"
 source = { editable = "packages/micro-agents-tools" }
 dependencies = [
     { name = "micro-agents" },