Skip to content

Python: prompt agents don't forward ActionRunContext to generate (tools/middleware miss request context) #5517

@jeffdh5

Description

@jeffdh5

How this was noticed

While wiring abort_signal through the prompt-agent generate path, we compared what _generate_prompt_agent_turn passes to generate_action against the registered /util/generate action handler in _generate.py.

abort_signal was forwarded; context was not.

The registered generate action explicitly pipes ActionRunContext.context into generate:

# py/packages/genkit/src/genkit/_ai/_generate.py — define_generate_action
async def generate_action_fn(input, ctx: ActionRunContext) -> ModelResponse:
    return await generate_with_request(
        registry=registry,
        raw_request=input,
        on_chunk=on_chunk,
        context=dict(ctx.context),          # ✅ forwarded
        abort_signal=ctx.abort_signal,
    )

Prompt-backed agents (define_agent / define_prompt_agent) call the same generate entrypoint but omit context:

# py/packages/genkit/src/genkit/_ai/_agent.py — _generate_prompt_agent_turn
response = await generate_action(
    child_registry,
    gen_options,
    on_chunk=_on_chunk,
    abort_signal=ctx.abort_signal,          # ✅ forwarded
    # context=... is missing               # ❌
)

That means tools invoked during a prompt-agent turn never receive caller context via GenerateMiddlewareContext.custom_context, even when the agent HTTP request was made with auth/metadata from a context_provider.

What we expect

Request context (auth, tenant id, etc.) should reach everything downstream of generate the same way it does for a direct /util/generate call or a custom agent that passes context= explicitly:

  1. HTTP context_provider → agent stream_bidi(..., context=...)
  2. Agent runtime → ActionRunContext.context
  3. Prompt turn → generate_action(..., context=dict(ctx.context))
  4. Generate → GenerateMiddlewareContext.custom_context
  5. Tools/models → ToolRunContext.context / model run(..., context=...)

Example: a tool that reads auth from context (see py/samples/context):

@ai.tool()
async def whoami(_: dict, ctx: ToolRunContext) -> str:
    user = ctx.context.get("auth", {}).get("uid")
    return f"hello {user}"

When this tool runs inside a prompt agent served over HTTP with a context provider, ctx.context is {} today. The same tool invoked via a flow or explicit generate(..., context=...) works.

JS reference: the agent runtime passes ambient context into the agent fn (context: getContext() in js/ai/src/agent.ts), and generate propagates context through its middleware ctx to model/tool runs. Python should match that product behavior via explicit custom_context plumbing.

Root cause (two layers)

1. Prompt turn doesn't forward context to generate

_generate_prompt_agent_turn is the only agent generate callsite that drops ctx.context.

2. Agent runtime may not populate ActionRunContext.context at all

SessionRunner.run builds a fresh context with only streaming + abort:

action_ctx = ActionRunContext(
    streaming_callback=self._send_chunk,
    abort_signal=abort_signal,
)

It does not seed from get_current_context() / the _action_context ContextVar that BidiAction.stream_bidi(..., context=...) sets. JS agent runtime explicitly passes context: getContext() when invoking the agent fn.

Additionally, the FastAPI agent HTTP handler currently calls agent.stream_bidi(init) without a context= argument, unlike flows which resolve context_provider and pass it to flow.run(..., context=...).

So even fixing _generate_prompt_agent_turn alone may not be enough until inbound HTTP context is wired into stream_bidi and SessionRunner.

Proposed fix

  1. SessionRunner.run: populate agent ActionRunContext from inbound context, e.g.

    from genkit._core._action import get_current_context
    
    action_ctx = ActionRunContext(
        streaming_callback=self._send_chunk,
        abort_signal=abort_signal,
        context=dict(get_current_context() or {}),
    )
  2. _generate_prompt_agent_turn: mirror the registered generate action:

    response = await generate_action(
        child_registry,
        gen_options,
        on_chunk=_on_chunk,
        context=dict(ctx.context),
        abort_signal=ctx.abort_signal,
    )
  3. FastAPI agent route (py/plugins/fastapi/handler.py): resolve context_provider(request) and pass to agent.stream_bidi(init, context=action_context) — same pattern as flow HTTP handling.

  4. Test: prompt agent + tool that asserts ctx.context["auth"] (or similar) is present when stream_bidi is invoked with a context dict; regression test that _generate_prompt_agent_turn forwards it into generate.

Scope note

This is separate from the recent abort_signal wiring work (that part is correct on the prompt path). Context and abort are independent channels on ActionRunContext; both should propagate through prompt agents.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    Status
    No status

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions