Skip to content

Returning new Content from on_user_message_callback fails to update invocation_context in runners.py #2909

@halfpasttense

Description

@halfpasttense

Returning a newly created Content object from on_user_message_callback plugin callback does not save the Content back into the invocation_context in runners.py:

async def _run_with_trace(
        new_message: types.Content,
    ) -> AsyncGenerator[Event, None]:
      with tracer.start_as_current_span('invocation'):
        session = await self.session_service.get_session(
            app_name=self.app_name, user_id=user_id, session_id=session_id
        )
        if not session:
          raise ValueError(f'Session not found: {session_id}')

        invocation_context = self._new_invocation_context(
            session,
            new_message=new_message,
            run_config=run_config,
        )
        root_agent = self.agent

        # Modify user message before execution.
        modified_user_message = await invocation_context.plugin_manager.run_on_user_message_callback(
            invocation_context=invocation_context, user_message=new_message
        )
        if modified_user_message is not None:
          new_message = modified_user_message

        if new_message:
          await self._append_new_message_to_session(
              session,
              new_message,
              invocation_context,
              run_config.save_input_blobs_as_artifacts,
              state_delta,
          )

        invocation_context.agent = self._find_agent_to_run(session, root_agent)

        async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
          async with Aclosing(ctx.agent.run_async(ctx)) as agen:
            async for event in agen:
              yield event

        async with Aclosing(
            self._exec_with_plugin(invocation_context, session, execute)
        ) as agen:
          async for event in agen:
            yield event

    async with Aclosing(_run_with_trace(new_message)) as agen:
      async for event in agen:
        yield event

Metadata

Metadata

Labels

core[Component] This issue is related to the core interface and implementationhelp wanted[Community] Extra attention is needed

Type

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions