feat: expose per-agent execution context

[thomwebb]

What

Adds a ContextVar-backed seam exposing the agent instance that owns the current model run:

• code_puppy/agent_execution_context.py — executing_agent_context(agent) + get_executing_agent()
• Scoped around the runtime's agent run in agents/_runtime.py (the task body owns its own scope, and agent_run_start/agent_run_end are wrapped where they fire) and around sub-agent invocation in tools/subagent_invocation.py. The innermost agent wins and outer state is restored on exit.

Coverage: run-scoped hooks — the agent task body (on_agent_run_context, tool-call hooks), agent_run_start/agent_run_end, and sub-agent runs. Not covered: build/registration-time hooks (on_register_agent_tools etc.) — those fire before a run exists and already receive agent_name explicitly, so they don't need this seam.

Why

Plugins currently have to consult the process-global agent manager to know "which agent is running", which is wrong under concurrent sessions and sub-agents. This gives them a run-scoped, async-safe answer. Needed by core-plugins for per-agent tool config (e.g. an agent that disables spill); the plugin side uses a guarded import, so it degrades gracefully on older runtimes and this lands independently.

Tests

• context nesting/restoration and cross-task isolation (tests/agents/test_execution_context.py)
• runtime agent task resolves its own agent, restored to None after
• sub-agent invocation scopes its own agent config

63 passed across the four touched test files; ruff clean.

Review follow-up (WSxDemise's review)

• The context now lives inside run_agent_task() (the body owns its scope, robust to scheduling changes) instead of wrapping the create_task call; agent_run_start/agent_run_end are wrapped where they fire — with a regression test asserting both hooks resolve the agent.
• get_executing_agent() is annotated BaseAgent | None under TYPE_CHECKING.
• Docstring notes the run_in_executor propagation caveat and why the module sits at top level.
• Test file renamed to tests/test_agent_execution_context.py per repo convention.
• The registration-time gap is addressed as a documented non-goal (above) rather than widened scope: registration hooks already get agent_name explicitly, and the motivating consumer (core-plugins spill) reads per-agent config at tool-result time, not registration time.
• The incidental import-sort hunks are lefthook (ruff --select I) artifacts, left as-is.

[WSxDemise]

[Wes's CodePuppy Agent Review]

Reviewed at head 5d63220d. CI green (quality / macOS 3.13 / windows-encoding). No prior human or bot reviews on this PR; qodo-merge[bot] is not configured for public GitHub, so there are no automated suggestions to triage.

Summary

Adds code_puppy/agent_execution_context.py — a 32-line ContextVar seam (executing_agent_context(agent) / get_executing_agent()) so plugins can resolve which agent instance is running instead of consulting the process-global get_current_agent(), which is wrong under concurrent ACP sessions and sub-agents.

Entered in two places:

• code_puppy/agents/_runtime.py:935 — wraps only the asyncio.create_task(...) call, relying on copy-at-task-creation.
• code_puppy/tools/subagent_invocation.py:457-460 — added to the existing subagent_context(...) block, scoping the whole sub-agent run.

The implementation is idiomatic and faithfully mirrors the established tools/subagent_context.py:90-140 pattern (ContextVar + @contextmanager + token reset in try/finally). Tests are meaningful rather than decorative. My concerns are about scope boundaries, not correctness — and findings 1–3 all collapse into a single edit.

---

Really Should Fix These

• [P1] The seam is dark at the exact hook the PR was written for — tool registration — code_puppy/agents/_builder.py:610,625, code_puppy/tools/subagent_invocation.py:422-425

  • Impact: The description states this is "Needed by core-plugins for per-agent tool config (e.g. an agent that disables spill)." But register_tools_for_agent() — which fires on_register_agent_tools(agent_name) — runs inside build_pydantic_agent() (_runtime.py:697/701) and at subagent_invocation.py:422-425, both before executing_agent_context(...) is entered. Same for on_wrap_pydantic_agent (_builder.py:636, subagent_invocation.py:428). A plugin calling get_executing_agent() from any of those gets None. The seam is only live from on_agent_run_context inward.
  • Fix: Either extend the scope to cover the build/wrap/registration phase in both paths, or amend the description to state precisely which hooks are covered and note that registration-time plugins should keep using the agent_name argument already passed to on_register_agent_tools.

• [P1] agent_run_start / agent_run_end cannot resolve the executing agent, and the two paths are asymmetric — code_puppy/agents/_runtime.py:908, code_puppy/agents/_runtime.py:1092

  • Impact: on_agent_run_start is awaited at line 908 and on_agent_run_end in the finally at line 1092 — both outside the with at line 935, which exits the moment create_task returns. These are the two most obvious "which agent is running" consumers in the codebase (callbacks.py:1039-1104 advertises them for token refresh, analytics, orchestration), and they are exactly the ones left blind. Meanwhile the sub-agent path wraps its entire block, so plugin authors get different behavior depending on which path they are on.
  • Fix: Widen the scope so it spans the run rather than task creation — see next item for the concrete shape.

• [P1] create_task-only scoping is fragile and will break silently under refactor — code_puppy/agents/_runtime.py:932-936

  • Impact: Correctness rests entirely on the implicit contextvars.copy_context() that asyncio.create_task performs; nothing at the call site enforces it. If line 936 later becomes await run_agent_task(), or the task creation moves, or it gets wrapped in a helper, the scope evaporates with no test failure at the point of the mistake and no runtime error — plugins just start seeing None. The three-line comment is doing all the load-bearing work.
  • Fix: Set the context inside run_agent_task() (_runtime.py:829) so the body owns its own scope unconditionally, or pass it explicitly via asyncio.create_task(coro, context=ctx). Placing it inside run_agent_task() and widening to enclose the on_agent_run_start/finally span resolves findings 1, 2, and 3 together.

---

Nits

• [P2] get_executing_agent() returns bare Any — code_puppy/agent_execution_context.py:30

  • Impact: Public, plugin-facing API (__all__, line 15) with no IDE completion or mypy protection for exactly the external consumers this PR targets.
  • Fix: if TYPE_CHECKING: from code_puppy.agents.base_agent import BaseAgent and annotate -> "BaseAgent | None". from __future__ import annotations (line 9) already makes this free at runtime.

• [P2] Docstring omits the thread-pool caveat — code_puppy/agent_execution_context.py:1-7,22

  • Impact: "work spawned within this async context" is not true for loop.run_in_executor(None, ...), which does not copy context (unlike asyncio.to_thread). This codebase does exactly that in tools/common.py:1412 and tools/command_runner.py:1477, so a plugin author will wrongly assume approval-backend and shell-executor threads can resolve the agent.
  • Fix: One line noting it propagates to tasks and to_thread, but not to run_in_executor workers. The sibling tools/common.py:113 documents this caveat — worth matching.

• [P2] Test file location doesn't mirror the module — tests/agents/test_execution_context.py

  • Impact: Module is top-level (code_puppy/), test is under tests/agents/. Convention for top-level modules is tests/test_<module>.py (e.g. code_puppy/skill_provider.py → tests/test_skill_provider.py).
  • Fix: Rename to tests/test_agent_execution_context.py.

• [P2] Module placement diverges from its closest sibling — code_puppy/agent_execution_context.py

  • Impact: The near-identical subagent_context.py lives under code_puppy/tools/, and this is an agents-runtime concern, so code_puppy/agents/ looks like the intuitive home. Top-level is nonetheless defensible: agents/__init__.py:20 imports run_stats for its registration side effects, so a plugin importing code_puppy.agents.* drags that in — top-level keeps the plugin import cheap and side-effect-free.
  • Fix: No change needed, but a one-line comment recording why it sits at top level would stop a future cleanup PR from "fixing" it into agents/ and reintroducing the side-effect import.

• [P3] Three unrelated import/whitespace hunks — _runtime.py:40, subagent_invocation.py:56-60, tests/test_subagent_invocation_usage.py:28-29

  • Impact: None functionally. These look auto-applied by lefthook.yml:14 (ruff check --select I --fix) rather than deliberate, but they are undescribed diff noise.
  • Fix: Leave them — reverting fights the repo's own pre-commit hook. Optionally mention the incidental import-sort in the description.

• [P3] No end-to-end nested main-run → sub-agent test — tests/agents/test_execution_context.py

  • Impact: Nesting is proven with two bare object() sentinels, and the two integration tests each cover one path in isolation. The realistic scenario — main agent running, sub-agent invoked mid-run, plugin asks "who's running?" — is never exercised end-to-end.
  • Fix: Optional follow-up asserting the sub-agent's agent_config wins inside the nested run and the main agent is restored afterward, locking in the "innermost agent wins" contract.

---

Scope & description parity

Scope is tight — no YAGNI (32 lines, no speculative parameters) and no DRY violation (checked every ContextVar( in code_puppy/; none already carries an agent instance). The only out-of-scope hunks are the three lint-churn lines above.

One parity gap: the "per-agent tool config" motivation in the description is not yet served by the code that landed (finding 1). Everything else in the description is backed by real code.

Nice work — this is a clean, well-tested change. Landing the scope widening would make the seam do what the description promises.

[thomwebb]

Thanks @WSxDemise — all three P1s addressed in c55d9fe: the scope now lives inside run_agent_task() (body owns it, scheduling-proof), and agent_run_start/agent_run_end are wrapped where they fire, with a new regression test asserting both resolve the agent. On the registration-time P1: taken as a documented non-goal rather than widened scope — registration hooks already receive agent_name explicitly, and the motivating consumer (spill's per-agent config) reads at tool-result time. Nits done: BaseAgent | None annotation under TYPE_CHECKING, executor-caveat docstring line, test file renamed to tests/test_agent_execution_context.py. Import-sort hunks are lefthook artifacts, left in place. 66 passed / ruff clean.

[WSxDemise]

[Wes's CodePuppyAgent Review]

Re-reviewed at c55d9fe9. Verified every change against the code rather than the summary. CI green (quality / macOS 3.13 / windows-encoding).

Verified

Item	Status	Evidence
P1 — scope owned by task body	Fixed	_runtime.py:829-836 — run_agent_task() wraps _run_agent_task_body(); create_task at :942 is now bare, so scheduling changes can't break it
P1 — agent_run_start / agent_run_end	Fixed	_runtime.py:917-922 and :1098-1110, wrapped where they fire
P1 — regression test	Present	test_agent_run_lifecycle_hooks_resolve_executing_agent asserts both hooks resolve the agent and it's None after; unregister_callback teardown confirmed at callbacks.py:240
P2 — BaseAgent | None	Done	agent_execution_context.py:38 under TYPE_CHECKING
P2 — executor caveat	Done	Module docstring, lines 9-10
P2 — top-level placement rationale	Done	Docstring lines 6-8
P2 — test rename	Done	tests/test_agent_execution_context.py
Description parity	Fixed	New Coverage/Not-covered paragraph + follow-up section now match the diff

Correction to my own review

My first P1 was over-called, and I want to be explicit about it. I claimed the seam was "dark at the exact hook the PR was written for" and that the motivating use case wasn't served. The factual half was right — registration-time hooks do fire before any scope exists. The impact half was wrong.

I inferred from the phrase "per-agent tool config" that the consumer read at registration time. I never checked. Core-plugins spill registers on post_tool_call (code_puppy_core_plugins/spill/register_callbacks.py:247), which fires at tool-result time — inside the task body, and therefore covered by the original 5d63220d scope, let alone this one. The package was installed in the venv the whole time; I could have confirmed it in one grep.

So your resolution is the correct one, and stronger than my proposed fix: widening scope to cover registration would have added surface area for a consumer that doesn't exist, when on_register_agent_tools already receives agent_name explicitly (callbacks.py:807). Documenting it as a non-goal is right. My apologies for the noise — I should have verified the consumer before asserting the use case was broken.

The other two P1s I'd still call fair: the create_task fragility was real regardless of consumer, and the lifecycle-hook gap was genuine.

Remaining

Nothing blocking. One optional follow-up, unchanged in priority from last round: no end-to-end test for a nested main-run → sub-agent invocation asserting the innermost agent wins and the outer is restored. The unit test covers nesting with sentinels and each path is covered in isolation, so this is a nice-to-have, not a gap.

on_agent_run_cancel in the sub-agent path (subagent_invocation.py:510) sits inside the executing_agent_context(agent_config) block — checked, correctly covered.

Approving. Clean fix, and the non-goal documentation is the better call.