Pass a per-run framework context/state through the Session to every framework adapter

.agents/skills/ak-dev-architecture/SKILL.md

@@ -34,6 +34,7 @@ Tracks state across related interactions. Key properties:
 - **Framework-specific data**: Stored via `get(key)` / `set(key, value)` — each framework stores its own state under a unique key (e.g., `"openai"`, `"langgraph"`)
 - **Volatile cache** (`v_cache`): Cleared after every `Runtime.run()` invocation — use for transient per-request data
 - **Non-volatile cache** (`nv_cache`): Persisted across requests within the session — use for data that should survive multiple interactions
+- **Reserved keys** (`Session.Keys` enum): `VOLATILE_CACHE`/`NON_VOLATILE_CACHE` back the two caches; `FRAMEWORK_CONTEXT` (`"framework_context"`) holds a per-run, framework-agnostic caller context/state **dict** (must be picklable) that runners inject into the native framework call and write back on success. Unlike the caches it is **not** pre-initialized — an unset key reads back as `None` (absent ⇒ no injection, no behaviour change). Callers seed/read it via `session.set/get(Session.Keys.FRAMEWORK_CONTEXT.value, ...)`
 - **Async context manager**: `async with session:` acquires a lock and sets the session as the current context via `contextvars`
 - **`Session.current()`**: Class method to retrieve the active session from any code running within the session context
 
@@ -55,6 +56,7 @@ Encapsulates framework-specific execution logic:
 - **`stream(agent, session, requests) -> AsyncGenerator[str, None]`**: Abstract async generator that yields token deltas for streaming execution (`execution.mode: stream`). Frameworks without native token streaming (CrewAI, smolagents) implement it by raising `NotImplementedError`
 - Each framework implements its own Runner (e.g., `OpenAIRunner`, `LangGraphRunner`, `CrewAIRunner`, `GoogleADKRunner`, `SmolagentsRunner`)
 - Runners handle: creating `ToolContext`, converting request models to framework-native formats, invoking the framework's execution API, converting responses back to `AgentReply`
+- **Per-run framework context**: the base `Runner` provides `_load_framework_context(session)` (returns a **deep copy** of the reserved `framework_context` key, or `None` when absent) and `_store_framework_context(session, incoming, produced)` (shallow-merges `produced` over `incoming` — framework-touched top-level keys win, untouched caller keys preserved — with a fail-fast picklability check). Each adapter's `run`/`stream` calls load before the native call, injects `incoming` via its native mechanism, and calls store **only after a successful native call** (inside the `try`, before the `except`; after the `async for` loop for streams, never in `finally`) so a crash/disconnect leaves the stored context intact. Round-trip fidelity is per-framework (OpenAI full; ADK all-but-internal-and-scope-prefixed keys, accumulate-only; smolagents pre-seeded keys only, and the context is also appended to the task prompt; LangGraph declared channels only; CrewAI unsupported — warns once per runner and skips). When an adapter seeds caller keys into a native state dict, write AK-internal keys **last** so a caller key cannot displace them (`ak_tool_context` in ADK, `messages` in LangGraph)
 
 ### Module (`ak-py/src/agentkernel/core/module.py`)
 

.agents/skills/ak-dev-new-framework-integration/SKILL.md

@@ -157,6 +157,52 @@ async def stream(self, agent: Any, session: Session, requests: list[AgentRequest
 
 `Runtime.stream()` wraps each yielded token in a `StreamChunk`, runs it through `PostHook.on_stream_chunk()`, and forwards it to the caller (REST SSE endpoint or AWS Lambda WebSocket/SQS pipeline). No other core changes are needed to support a new framework's streaming — just implement `Runner.stream()`.
 
+### 3c. Wire up the per-run framework context
+
+The base `Runner` provides two helpers so a caller-supplied, framework-agnostic context/state dict
+(the reserved `Session.Keys.FRAMEWORK_CONTEXT` key) rides across turns. **Your `run()` and
+`stream()` must call them and map the one AK-level dict onto your framework's native
+context/state mechanism** (or decline it explicitly, as CrewAI does):
+
+- `incoming = self._load_framework_context(session)` — call **before** the native invocation.
+  Returns a **deep copy** of the stored dict, or `None` when the key is absent. When `None`, inject
+  nothing (framework default) — this keeps the no-context path unchanged for existing apps.
+- Inject `incoming` (when not `None`) via the framework's native mechanism (a run `context=`, an
+  input state channel, a session-state delta, `additional_args=`, …).
+- After a **successful** native call, extract the framework's post-run state as `produced` and call
+  `self._store_framework_context(session, incoming, produced)`. This shallow-merges `produced` over
+  `incoming` (framework-touched top-level keys win; untouched caller keys preserved) and fail-fast
+  checks picklability before writing back.
+
+```python
+# In run(), inside the existing try, around the native call:
+incoming = self._load_framework_context(session)
+result = await self._execute(agent, fw_session, prompt, context=incoming)  # inject natively
+produced = self._extract_state(result, incoming)  # framework-specific; may be a subset of keys
+self._store_framework_context(session, incoming, produced)  # only after a successful call
+```
+
+**Placement matters (atomicity):** put the write-back **inside the `try`, after the native call,
+before the `except`** — and for `stream()`, **after the `async for` loop but still inside the
+`try`**, never in `finally`. A framework error or a client disconnect (`GeneratorExit`) then unwinds
+before it, leaving the previously stored context intact rather than persisting partial state.
+
+**Seed AK-internal keys last.** If you inject the caller's dict by merging it into a native state
+dict that also carries AK-internal entries, assign the internal ones **after** the caller's keys so a
+caller key can never displace them (`ak_tool_context` in ADK, `messages` in LangGraph). The failures
+this prevents are silent and confusing — a broken tool-context lookup, or a replaced message list.
+
+**Watch for injection side effects.** A framework's "context" slot is not always private: smolagents'
+`additional_args` is merged into the agent state *and* appended to the task prompt, so the caller's
+dict reaches the model. If your framework does something similar, document it on the framework's page
+so callers know not to put secrets in `framework_context`.
+
+**Declare your round-trip fidelity honestly** in the framework's docs and the fidelity table in
+`docs/docs/core-concepts/runner.md` — how much of a caller dict actually survives depends on the
+framework (full round-trip, filtered to seeded keys, declared-channels-only, or unsupported). If the
+framework has no safe caller-state slot, do **not** inject; instead log a single warning per run and
+skip both load and write-back (see the CrewAI adapter for the pattern).
+
 ### 4. Implement the Agent Wrapper
 
 Subclass `Agent` from `agentkernel.core.base`:

.github/test-config.yaml

@@ -12,12 +12,18 @@ e2e:
       path: examples/cli/openai-dynamic
     - type: cli
       path: examples/cli/openai_structured
+    - type: cli
+      path: examples/cli/openai_context
     - type: cli
       path: examples/cli/adk
+    - type: cli
+      path: examples/cli/adk_context
     - type: cli
       path: examples/cli/crewai
     - type: cli
       path: examples/cli/langgraph
+    - type: cli
+      path: examples/cli/langgraph_context
     - type: cli
       path: examples/cli/multi
     - type: cli

ak-py/src/agentkernel/cli/cli.py

@@ -1,9 +1,8 @@
 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 +46,33 @@ def list(self):
                 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 runs on a daemon thread so Ctrl+C exits promptly 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 @@ async def run(self):
         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("!"):

ak-py/src/agentkernel/core/base.py

@@ -1,15 +1,32 @@
 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
 
 from .hooks import PostHook, PreHook
 from .model import AgentReply, AgentRequest
 from .util.key_value_cache import KeyValueCache
 
+_log = logging.getLogger("ak.core.runner")
+
+
+def _not_picklable(value: Any) -> bool:
+    """
+    Checks whether the given value can be pickled.
+    :param value: The value to test for picklability.
+    :return: True if the value cannot be pickled, otherwise False.
+    """
+    try:
+        pickle.dumps(value)
+        return False
+    except Exception:
+        return True
+
 
 class Session:
     """
@@ -39,6 +56,7 @@ class Keys(Enum):
 
         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)
 
@@ -220,6 +238,83 @@ def name(self) -> str:
         """
         return self._name
 
+    def _load_framework_context(self, session: Session | None) -> dict | None:
+        """
+        Returns a deep copy of the per-run framework context stored in the session, or None if the key is absent.
+        The copy isolates the stored value from in-run mutation, so a failed run leaves the previously stored
+        context intact.
+        :param session: The session to read the reserved framework_context key from, or None.
+        :return: A deep copy of the stored dict, or None when the key is absent.
+        :raises TypeError: If the stored value is not a dict.
+        """
+        if session is None:
+            return None
+        stored = session.get(Session.Keys.FRAMEWORK_CONTEXT.value)
+        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:
+        """
+        Shallow-merges the framework's post-run state over the context loaded this turn and writes the result
+        back to the reserved framework_context key. Keys the framework touched win; untouched caller keys are
+        preserved. No-op when the key was absent this turn.
+        :param session: The session to write the merged framework_context back to, or None.
+        :param incoming: The deep copy loaded this turn, or None when the key was absent.
+        :param produced: The framework's post-run state delta, 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(Session.Keys.FRAMEWORK_CONTEXT.value, merged)
+
+    @staticmethod
+    def _ensure_framework_context_picklable(session: Session, ctx: Mapping[str, Any]) -> None:
+        """
+        Fails fast if the merged context cannot be pickled. Sessions are persisted with pickle, so a
+        non-picklable value would otherwise abort the whole session store() with an opaque error.
+        :param session: The session whose id is named in the error.
+        :param ctx: The merged 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 _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 streamed-run framework_context write-back failure instead of raising it. Raising from a stream
+        would escape the generator as a transport error and skip Runtime.stream's session store(), losing the
+        whole turn's state rather than just the context.
+        :param session: The session whose write-back failed, named in the log message.
+        :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 +446,18 @@ def _attach_system_tools(self) -> None:
 
         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-native agent's list-based `tools` attribute,
+        creating the list when it is absent and skipping tools that are already attached. Shared by
+        the framework agents whose native tool collection is a plain list.
+        :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)