agentkernel 0.8.0__tar.gz → 0.8.1__tar.gz

{agentkernel-0.8.0 → agentkernel-0.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agentkernel
-Version: 0.8.0
+Version: 0.8.1
 Summary: Agent Kernel - Unified AI Agents Runtime
 Author: Yaala Labs
 Author-email: Yaala Labs <agentkernel@yaalalabs.com>

{agentkernel-0.8.0 → agentkernel-0.8.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "agentkernel"
-version = "0.8.0"
+version = "0.8.1"
 description = "Agent Kernel - Unified AI Agents Runtime"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/cli/cli.py

@@ -1,9 +1,9 @@
 import asyncio
 import logging
 import readline  # Enables line editing and history features for input() in the CLI
+import threading
 
 from ..core import AgentService
-from ..core.config import AKConfig
 
 ak_cli_logger = logging.getLogger("ak.cli")
 
@@ -47,6 +47,32 @@ class CLI:
                 self._print(f"  {agent.name}")
             self._print()
 
+    @staticmethod
+    async def _ainput(prompt: str) -> str:
+        """
+        Reads a line of input without blocking the event loop, so background tasks keep running between turns.
+        The reader is a daemon thread, so Ctrl+C exits instead of waiting on a thread parked in `input()`.
+        :param prompt: The prompt to display to the user.
+        :return: The line entered by the user.
+        """
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[str] = loop.create_future()
+
+        def resolve(setter, value):
+            if not future.done():
+                setter(value)
+
+        def read():
+            try:
+                line = input(prompt)
+            except BaseException as e:  # noqa: BLE001 - relayed to the awaiting caller
+                loop.call_soon_threadsafe(resolve, future.set_exception, e)
+            else:
+                loop.call_soon_threadsafe(resolve, future.set_result, line)
+
+        threading.Thread(target=read, daemon=True, name="ak-cli-input").start()
+        return await future
+
     async def run(self):
         self._print("AgentKernel CLI (type !help for commands or !quit to exit):")
         self._service.select()
@@ -57,7 +83,7 @@ class CLI:
         while True:
             try:
                 name = self._service.agent.name if self._service.agent else "none"
-                prompt = input(f"({name}) >> ")
+                prompt = await self._ainput(f"({name}) >> ")
                 if not prompt.strip():
                     continue
                 if prompt.startswith("!"):

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/core/base.py

@@ -1,8 +1,10 @@
 import asyncio
 import contextvars
+import copy
 import logging
+import pickle
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, Iterator
+from collections.abc import AsyncGenerator, Iterator, Mapping
 from enum import Enum
 from typing import Any, ClassVar, Self, cast
 
@@ -10,6 +12,8 @@ from .hooks import PostHook, PreHook
 from .model import AgentReply, AgentRequest
 from .util.key_value_cache import KeyValueCache
 
+_log = logging.getLogger("ak.core.runner")
+
 
 class Session:
     """
@@ -39,6 +43,7 @@ class Session:
 
         VOLATILE_CACHE = "v_cache"
         NON_VOLATILE_CACHE = "nv_cache"
+        FRAMEWORK_CONTEXT = "framework_context"
 
     current_session: ClassVar[contextvars.ContextVar[Self | None]] = contextvars.ContextVar("current_session", default=None)
 
@@ -157,6 +162,35 @@ class Session:
         """
         return cast(KeyValueCache, self.get(Session.Keys.NON_VOLATILE_CACHE.value))
 
+    def get_framework_context(self) -> dict | None:
+        """
+        Returns the live per-run framework context dict, so hooks can edit it in place.
+        Never auto-creates the key: absent means "nothing injected", {} means "injected but empty".
+        :return: The stored framework context dict, or None when the key is absent.
+        """
+        return self.get(Session.Keys.FRAMEWORK_CONTEXT.value)
+
+    def set_framework_context(self, context: dict) -> dict:
+        """
+        Seeds or replaces the per-run framework context carried across turns.
+        :param context: The context dict to store. Must be picklable, since sessions are persisted with pickle.
+        :return: The stored context dict.
+        :raises TypeError: If context is not a dict.
+        """
+        if not isinstance(context, dict):
+            raise TypeError(
+                f"Session '{self._id}' framework_context must be a dict, got "
+                f"{type(context).__name__}. The reserved framework_context key carries a "
+                f"per-run context/state dict; wrap the value in a dict before setting it."
+            )
+        return self.set(Session.Keys.FRAMEWORK_CONTEXT.value, context)
+
+    def clear_framework_context(self) -> None:
+        """
+        Deletes the per-run framework context, so nothing is injected on the next turn.
+        """
+        self.delete(Session.Keys.FRAMEWORK_CONTEXT.value)
+
     def set(self, key: str, value: Any) -> Any:
         """
         Sets a session data object for the specified key.
@@ -220,6 +254,92 @@ class Runner(ABC):
         """
         return self._name
 
+    def _load_framework_context(self, session: Session | None) -> dict | None:
+        """
+        Loads the per-run framework context to inject into this run.
+        Returns a deep copy, so a failed run leaves the stored context intact.
+        :param session: The session to read the framework context from, or None.
+        :return: A deep copy of the stored dict, or None when nothing is stored.
+        :raises TypeError: If the stored value is not a dict.
+        """
+        if session is None:
+            return None
+        stored = session.get_framework_context()
+        if stored is None:
+            return None
+        if not isinstance(stored, dict):
+            raise TypeError(
+                f"Session '{session.id}' framework_context must be a dict, got "
+                f"{type(stored).__name__}. The reserved framework_context key carries a "
+                f"per-run context/state dict; wrap the value in a dict before setting it."
+            )
+        return copy.deepcopy(stored)
+
+    def _store_framework_context(self, session: Session | None, incoming: dict | None, produced: Mapping[str, Any] | None) -> None:
+        """
+        Merges the framework's post-run state over the context loaded this turn and stores the result.
+        Keys the framework touched win; untouched caller keys are preserved. No-op when nothing was loaded.
+        :param session: The session to write the merged context back to, or None.
+        :param incoming: The context loaded this turn, or None when nothing was stored.
+        :param produced: The framework's post-run state, or None.
+        """
+        if session is None or incoming is None:
+            return
+        merged = dict(incoming)
+        if produced:
+            merged.update(produced)
+        self._ensure_framework_context_picklable(session, merged)
+        session.set_framework_context(merged)
+
+    @staticmethod
+    def _not_picklable(value: Any) -> bool:
+        """
+        Checks whether the given value can be pickled.
+        :param value: The value to test.
+        :return: True if the value cannot be pickled.
+        """
+        try:
+            pickle.dumps(value)
+            return False
+        except Exception:
+            return True
+
+    @classmethod
+    def _ensure_framework_context_picklable(cls, session: Session, ctx: Mapping[str, Any]) -> None:
+        """
+        Fails fast if the context cannot be pickled, before it can break the session store.
+        :param session: The session whose id is named in the error.
+        :param ctx: The framework context to validate.
+        :raises TypeError: If any value in ctx is not pickle-serializable.
+        """
+        try:
+            pickle.dumps(ctx)
+        except Exception as exc:
+            offender = next(
+                (f"{k!r} ({type(v).__name__})" for k, v in ctx.items() if cls._not_picklable(v)),
+                "<unknown key>",
+            )
+            raise TypeError(
+                f"Session '{session.id}' framework_context is not picklable; "
+                f"offending entry: {offender}. framework_context values must be "
+                f"pickle-serializable so the session can be persisted."
+            ) from exc
+
+    @staticmethod
+    def _log_framework_context_stream_failure(session: Session | None, error: Exception) -> None:
+        """
+        Logs a framework context write-back failure at the end of a streamed run instead of raising it.
+        Raising here would surface as a transport error and skip the turn's session store.
+        :param session: The session whose write-back failed.
+        :param error: The exception raised while producing or storing the context.
+        """
+        session_id = session.id if session is not None else "<none>"
+        _log.error(
+            f"Session '{session_id}': framework_context write-back was skipped at the end of a "
+            f"streamed run; the previously stored context is left intact. Error: {error}",
+            exc_info=error,
+        )
+
     @abstractmethod
     async def run(self, agent: Any, session: Session, requests: list[AgentRequest]) -> AgentReply:
         """
@@ -351,3 +471,16 @@ class Agent(ABC):
 
         for tool in SystemToolFactory.get_all(self.name):
             self.attach_tool(tool.func)
+
+    @staticmethod
+    def _append_tools(agent: Any, wrapped: list[Any]) -> None:
+        """
+        Appends already-wrapped tools to a framework agent's list-based `tools` attribute, skipping duplicates.
+        :param agent: The framework-native agent holding the `tools` attribute.
+        :param wrapped: The wrapped tools to attach.
+        """
+        for w in wrapped:
+            if not hasattr(agent, "tools") or agent.tools is None:
+                agent.tools = []
+            if w not in agent.tools:
+                agent.tools.append(w)

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/framework/adk/adk.py

@@ -13,7 +13,7 @@ from google.adk.agents import BaseAgent
 from google.adk.agents.run_config import RunConfig, StreamingMode
 from google.adk.events import Event, EventActions
 from google.adk.runners import Runner
-from google.adk.sessions import BaseSessionService, InMemorySessionService
+from google.adk.sessions import BaseSessionService, InMemorySessionService, State
 from google.adk.tools import FunctionTool, ToolContext
 from google.genai import types
 from pydantic import ValidationError
@@ -76,6 +76,21 @@ class GoogleADKSession:
             )
             await self._session_service.append_event(self._session, event)
 
+    async def get_state(self) -> dict:
+        """
+        Returns the caller-visible ADK session state, stripping the AK-internal ``ak_tool_context`` id and the
+        ``app:`` / ``user:`` / ``temp:`` prefixed keys that are not session-scoped.
+        :return: The accumulated session state, or an empty dict when no session exists.
+        """
+        if self._session is None:
+            return {}
+        refreshed = await self._session_service.get_session(
+            app_name=self._session.app_name, user_id=self._session.user_id, session_id=self._session.id
+        )
+        state = dict(getattr(refreshed, "state", {}) or {})
+        state.pop("ak_tool_context", None)
+        return {k: v for k, v in state.items() if not k.startswith((State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX))}
+
 
 class GoogleADKRunner(BaseRunner):
     def __init__(self):
@@ -131,7 +146,7 @@ class GoogleADKRunner(BaseRunner):
                     parts.append(types.Part(file_data=types.FileData(file_uri=base64_data)))
                     continue
 
-                if base64_data.startswith(("data:")):
+                if base64_data.startswith("data:"):
                     mime_type = base64_data.split(";")[0][5:]
                 else:
                     if not req.mime_type:
@@ -143,15 +158,19 @@ class GoogleADKRunner(BaseRunner):
 
         return prompt, parts
 
-    async def _setup_session_context(self, agent: Any, session: Session, requests: list[AgentRequest]) -> tuple[str, Runner, AKToolContext]:
+    async def _setup_session_context(
+        self, agent: Any, session: Session, requests: list[AgentRequest], injected: dict | None = None
+    ) -> tuple[str, Runner, AKToolContext, GoogleADKSession]:
         """
         Setup ADK session and tool context.
         :param agent: The ADK agent.
         :param session: The AgentKernel session.
         :param requests: The requests.
-        :return: Tuple of (user_id, runner, tool_context). The caller is responsible for entering/exiting
-                 the returned tool_context around the runner's actual execution, since tools invoked by
-                 the agent look up this context by id from the cache while the agent is running.
+        :param injected: The per-run framework context to seed into the ADK session state, or None.
+        :return: Tuple of (user_id, runner, tool_context, adk_session). The caller is responsible for
+                 entering/exiting the returned tool_context around the runner's actual execution, since
+                 tools invoked by the agent look up this context by id from the cache while the agent is
+                 running. The returned adk_session lets the caller read state back without re-fetching.
         """
         app_name = "AgentKernel"
         user_id = "AgentKernel"
@@ -159,10 +178,13 @@ class GoogleADKRunner(BaseRunner):
 
         ctx: AKToolContext = AKToolContext(Runtime.current(), agent, session, requests)
         await adk_session.create_session(app_name=app_name, user_id=user_id, session_id=session.id)
-        await adk_session.update_session_state(ctx.id, agent.name, {"ak_tool_context": ctx.id})
+        # Assigned last so a caller key of the same name cannot replace the id tools resolve their context by.
+        state = dict(injected or {})
+        state["ak_tool_context"] = ctx.id
+        await adk_session.update_session_state(ctx.id, agent.name, state)
 
         runner = Runner(agent=agent.agent, app_name=app_name, session_service=adk_session.session_service)
-        return user_id, runner, ctx
+        return user_id, runner, ctx, adk_session
 
     @staticmethod
     async def get_response(runner: Runner, user_id: str, session_id: str, parts: list[types.Part]) -> str:
@@ -178,15 +200,12 @@ class GoogleADKRunner(BaseRunner):
         response_text = ""
 
         if hasattr(runner, "run_async"):
-            events = runner.run_async(user_id=user_id, session_id=session_id, new_message=new_message)
-            try:
-                async for event in events:
-                    if event.is_final_response() and event.content and event.content.parts:
-                        text_parts = [p.text for p in event.content.parts if hasattr(p, "text") and p.text]
-                        response_text = " ".join(text_parts) if text_parts else ""
-                        break
-            finally:
-                await events.aclose()
+            # Drain the stream instead of breaking early: stopping early cancels ADK's still-running root agent
+            # task, and with sub-agents the last final response is the one to return.
+            async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=new_message):
+                if event.is_final_response() and event.content and event.content.parts:
+                    text_parts = [p.text for p in event.content.parts if hasattr(p, "text") and p.text]
+                    response_text = " ".join(text_parts) if text_parts else ""
         else:
             for event in runner.run(user_id=user_id, session_id=session_id, new_message=new_message):
                 if event.is_final_response() and event.content and event.content.parts:
@@ -203,16 +222,24 @@ class GoogleADKRunner(BaseRunner):
         :param requests: The requests to the agent.
         :return: The result of the agent's execution.
         """
+        prompt = ""
         try:
             prompt, parts = self._process_requests(requests)
 
             if not parts:
                 return AgentReplyText(response="Sorry. No valid content found in the requests")
 
-            user_id, runner, ctx = await self._setup_session_context(agent, session, requests)
+            incoming = self._load_framework_context(session)
+            user_id, runner, ctx, adk_session = await self._setup_session_context(agent, session, requests, incoming)
             with ctx:
                 reply = await self.get_response(runner=runner, session_id=session.id, parts=parts, user_id=user_id)
 
+            # Write back the full ADK state, so keys a tool added during the run also round-trip. Done after
+            # the run, so a framework error leaves the stored context intact.
+            if incoming is not None:
+                produced = await adk_session.get_state()
+                self._store_framework_context(session, incoming, produced)
+
             output_schema = getattr(agent.agent, "output_schema", None)
             if output_schema is not None:
                 try:
@@ -237,7 +264,8 @@ class GoogleADKRunner(BaseRunner):
         if not parts:
             return
 
-        user_id, runner, ctx = await self._setup_session_context(agent, session, requests)
+        incoming = self._load_framework_context(session)
+        user_id, runner, ctx, adk_session = await self._setup_session_context(agent, session, requests, incoming)
         new_message = types.Content(role="user", parts=parts)
         run_config = RunConfig(streaming_mode=StreamingMode.SSE)
 
@@ -257,6 +285,15 @@ class GoogleADKRunner(BaseRunner):
                     if chunk:
                         yield chunk
 
+                # Only after the stream drains normally, so a disconnect or mid-stream error leaves the stored
+                # context intact. Deliberately not in a finally.
+                if incoming is not None:
+                    try:
+                        produced = await adk_session.get_state()
+                        self._store_framework_context(session, incoming, produced)
+                    except Exception as e:
+                        self._log_framework_context_stream_failure(session, e)
+
 
 class GoogleADKAgent(AKBaseAgent):
     """
@@ -293,9 +330,8 @@ class GoogleADKAgent(AKBaseAgent):
         Appends the given prompt text to the ADK agent's description.
         Called by the base Agent._setup_system_prompt() at init when multimodal is enabled.
         """
-        if hasattr(self._agent, "description") and self._agent.description:
-            if prompt not in self._agent.description:
-                self._agent.description += "\n" + prompt
+        if hasattr(self._agent, "description") and self._agent.description and prompt not in self._agent.description:
+            self._agent.description += "\n" + prompt
 
     def attach_tool(self, tool: Any) -> None:
         """
@@ -304,12 +340,7 @@ class GoogleADKAgent(AKBaseAgent):
         :param tool: Raw Python callable or already-wrapped ADK FunctionTool.
         """
         # Delegate to the tool builder to handle binding
-        wrapped = GoogleADKToolBuilder.bind([tool])
-        for w in wrapped:
-            if not hasattr(self._agent, "tools") or self._agent.tools is None:
-                self._agent.tools = []
-            if w not in self._agent.tools:
-                self._agent.tools.append(w)
+        self._append_tools(self._agent, GoogleADKToolBuilder.bind([tool]))
 
     def get_a2a_card(self):
         """
@@ -410,15 +441,40 @@ class GoogleADKToolBuilder(ToolBuilder):
         if not callable(func):
             raise TypeError(f"Expected a callable, got {type(func).__name__}")
 
+        signature = inspect.signature(func)
+        parameters = list(signature.parameters.values())
+
+        # ADK-aware tools declare `tool_context` themselves and expect the live ADK
+        # context to reach them, so it has to be forwarded. Generic tools do not, and
+        # for those the parameter is consumed here purely to activate the AK context.
+        declares_tool_context = "tool_context" in signature.parameters
+        tool_context_position = list(signature.parameters).index("tool_context") if declares_tool_context else -1
+
+        def forwarded_kwargs(args: tuple[Any, ...], kwargs: dict[str, Any], tool_context: ToolContext | None) -> dict[str, Any]:
+            """
+            Builds the keyword arguments for the wrapped function.
+
+            :param args: Positional arguments the wrapper was called with.
+            :param kwargs: Keyword arguments the wrapper was called with.
+            :param tool_context: The ADK tool context supplied to the wrapper.
+            :return: The keyword arguments to pass to the wrapped function.
+            """
+            if not declares_tool_context or "tool_context" in kwargs:
+                return kwargs
+            # The caller already bound tool_context positionally; do not bind it twice.
+            if len(args) > tool_context_position:
+                return kwargs
+            return {**kwargs, "tool_context": tool_context}
+
         if asyncio.iscoroutinefunction(func):
 
             @functools.wraps(func)
-            async def wrapper(*args: Any, tool_context: ToolContext, **kwargs: Any) -> Any:
+            async def wrapper(*args: Any, tool_context: ToolContext | None = None, **kwargs: Any) -> Any:
                 tctx: AKToolContext | None = None
                 try:
                     if tool_context and tool_context.state and tool_context.state.get("ak_tool_context"):
                         tctx = AKToolContext.fetch(tool_context.state["ak_tool_context"]).set()
-                    return await func(*args, **kwargs)
+                    return await func(*args, **forwarded_kwargs(args, kwargs, tool_context))
                 finally:
                     if tctx:
                         tctx.reset()
@@ -426,23 +482,20 @@ class GoogleADKToolBuilder(ToolBuilder):
         else:
 
             @functools.wraps(func)
-            def wrapper(*args: Any, tool_context: ToolContext, **kwargs: Any) -> Any:
+            def wrapper(*args: Any, tool_context: ToolContext | None = None, **kwargs: Any) -> Any:
                 tctx: AKToolContext | None = None
                 try:
                     if tool_context and tool_context.state and tool_context.state.get("ak_tool_context"):
                         tctx = AKToolContext.fetch(tool_context.state["ak_tool_context"]).set()
-                    return func(*args, **kwargs)
+                    return func(*args, **forwarded_kwargs(args, kwargs, tool_context))
                 finally:
                     if tctx:
                         tctx.reset()
 
-        signature = inspect.signature(func)
-        parameters = list(signature.parameters.values())
-
         # Only add a tool_context parameter if the original function does not already
         # declare one, and insert it before any **kwargs (VAR_KEYWORD) parameter to
         # maintain a valid inspect.Signature ordering.
-        if "tool_context" not in signature.parameters:
+        if not declares_tool_context:
             tool_context_param = inspect.Parameter(
                 "tool_context",
                 kind=inspect.Parameter.KEYWORD_ONLY,

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/framework/crewai/crewai.py

@@ -273,6 +273,8 @@ class CrewAIRunner(Runner):
         """
         super().__init__(FRAMEWORK)
         self._log = logging.getLogger("ak.crewai.runner")
+        self._context_warned = False
+        """Whether the unsupported-framework_context warning was already logged."""
 
     def _transcript(self, session: Session) -> list[str] | None:
         """
@@ -375,6 +377,11 @@ class CrewAIRunner(Runner):
                 verbose=False,
                 memory=memory,
             )
+            # CrewAI's kickoff(inputs=...) are template-interpolation variables, not a context/state object, so
+            # there is no per-run caller-state slot. Warn once, and leave the stored context untouched.
+            if not self._context_warned and session is not None and session.get_framework_context():
+                self._log.warning("framework_context is set but CrewAI does not support per-run caller context/state; ignoring it.")
+                self._context_warned = True
             reply = await crew.kickoff_async(inputs={})
             if isinstance(getattr(reply, "pydantic", None), BaseModel):
                 agent_reply: AgentReply = AgentReplyAny(content=reply.pydantic.model_dump(mode="json"), prompt=prompt)
@@ -510,12 +517,7 @@ class CrewAIAgent(BaseAgent):
         :param tool: Raw Python callable or already-wrapped CrewAI tool.
         """
         # Delegate to the tool builder to handle binding
-        wrapped = CrewAIToolBuilder.bind([tool])
-        for w in wrapped:
-            if not hasattr(self.agent, "tools") or self.agent.tools is None:
-                self.agent.tools = []
-            if w not in self.agent.tools:
-                self.agent.tools.append(w)
+        self._append_tools(self.agent, CrewAIToolBuilder.bind([tool]))
 
     def override_system_prompt(self, prompt: str) -> None:
         """

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/framework/langgraph/langgraph.py

@@ -371,6 +371,7 @@ class LangGraphRunner(BaseRunner):
         :param requests: The requests to the agent.
         :return: The result of the agent's execution.
         """
+        prompt = ""
         context: ToolContext | None = None
         try:
             context = ToolContext(Runtime.current(), agent, session, requests).set()
@@ -387,10 +388,24 @@ class LangGraphRunner(BaseRunner):
 
             config, messages = self._prepare_session_and_messages(agent, session, prompt)
 
+            # Spread the context's top-level keys into the input state so they map onto the graph's state
+            # channels. `messages` is written last so a caller key cannot replace it.
+            incoming = self._load_framework_context(session)
+            input_state: dict[str, Any] = {}
+            if incoming:
+                input_state.update(incoming)
+            input_state["messages"] = messages
+
             result = await agent.agent.ainvoke(
-                input={"messages": messages},
+                input=input_state,
                 config=config,
             )
+
+            # Only keys the graph declares as state channels come back on `result`; the rest keep their value.
+            if incoming is not None:
+                produced = {k: result[k] for k in incoming if k in result}
+                self._store_framework_context(session, incoming, produced)
+
             structured = AgentReplyAny.from_output(result.get("structured_response"), prompt)
             if structured is not None:
                 return structured
@@ -423,8 +438,14 @@ class LangGraphRunner(BaseRunner):
 
             config, messages = self._prepare_session_and_messages(agent, session, prompt)
 
+            incoming = self._load_framework_context(session)
+            input_state: dict[str, Any] = {}
+            if incoming:
+                input_state.update(incoming)
+            input_state["messages"] = messages
+
             async for event in agent.agent.astream_events(
-                input={"messages": messages},
+                input=input_state,
                 config=config,
                 version="v2",
             ):
@@ -436,6 +457,16 @@ class LangGraphRunner(BaseRunner):
                         for item in content:
                             if isinstance(item, dict) and item.get("text"):
                                 yield item["text"]
+
+            # astream_events yields events, not a final state, so read the state back once the stream drains
+            # normally. A disconnect or mid-stream error unwinds first, leaving the stored context intact.
+            if incoming is not None:
+                try:
+                    state = await agent.agent.aget_state(config)
+                    produced = {k: state.values[k] for k in incoming if k in state.values}
+                    self._store_framework_context(session, incoming, produced)
+                except Exception as e:
+                    self._log_framework_context_stream_failure(session, e)
         finally:
             if context is not None:
                 context.reset()

{agentkernel-0.8.0 → agentkernel-0.8.1}/src/agentkernel/framework/openai/openai.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 from collections.abc import AsyncGenerator
 from typing import Any, Callable, List
 
@@ -171,6 +172,7 @@ class OpenAIRunner(BaseRunner):
         :param requests: The requests to the agent.
         :return: The result of the agent's execution.
         """
+        prompt = ""
         context: ToolContext | None = None
         try:
             context = ToolContext(Runtime.current(), agent, session, requests).set()
@@ -180,7 +182,13 @@ class OpenAIRunner(BaseRunner):
                 return AgentReplyText(response="Sorry. No valid content found in the requests")
 
             input_data, session_to_use = self._get_run_input(agent, session, prompt, message_content)
-            reply = (await Runner.run(agent.agent, input_data, session=session_to_use)).final_output
+            # Injected as the run context, so tools read and write it via RunContextWrapper.context.
+            # A deep copy is passed in so tools mutating it in place don't also mutate `incoming`.
+            incoming = self._load_framework_context(session)
+            produced = copy.deepcopy(incoming)
+            reply = (await Runner.run(agent.agent, input_data, session=session_to_use, context=produced)).final_output
+
+            self._store_framework_context(session, incoming, produced)
 
             structured = AgentReplyAny.from_output(reply, prompt)
             if structured is not None:
@@ -211,12 +219,21 @@ class OpenAIRunner(BaseRunner):
                 return
 
             input_data, session_to_use = self._get_run_input(agent, session, prompt, message_content)
-            result = Runner.run_streamed(agent.agent, input_data, session=session_to_use)
+            incoming = self._load_framework_context(session)
+            produced = copy.deepcopy(incoming)
+            result = Runner.run_streamed(agent.agent, input_data, session=session_to_use, context=produced)
 
             async for event in result.stream_events():
                 if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
                     if event.data.delta:
                         yield event.data.delta
+
+            # Only after the stream drains normally, so a disconnect or framework error leaves the stored
+            # context intact. Deliberately not in a finally.
+            try:
+                self._store_framework_context(session, incoming, produced)
+            except Exception as e:
+                self._log_framework_context_stream_failure(session, e)
         finally:
             if context is not None:
                 context.reset()
@@ -268,12 +285,7 @@ class OpenAIAgent(BaseAgent):
         :param tool: Raw Python callable or already-wrapped OpenAI function_tool.
         """
         # Delegate to the tool builder to handle binding
-        wrapped = OpenAIToolBuilder.bind([tool])
-        for w in wrapped:
-            if not hasattr(self._agent, "tools") or self._agent.tools is None:
-                self._agent.tools = []
-            if w not in self._agent.tools:
-                self._agent.tools.append(w)
+        self._append_tools(self._agent, OpenAIToolBuilder.bind([tool]))
 
     def get_a2a_card(self):
         """