state_delta is silently ignored when resuming an invocation without a new_message #6644

"""state_delta when resuming an invocation without a new_message.

Covers both dispatch paths in Runner.run_async:
  A) node path   -> _run_node_async                        (LlmAgent, a BaseNode)
  B) legacy path -> _setup_context_for_resumed_invocation  (plain BaseAgent)
"""

import asyncio
from typing import AsyncGenerator

from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import LlmAgent
from google.adk.apps.app import App, ResumabilityConfig
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.events.event import Event
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types

USER, SESSION = "u1", "s1"


class EchoLlm(BaseLlm):
  """Minimal offline model so the repro needs no network or API key."""

  model: str = "echo"

  async def generate_content_async(
      self, llm_request, stream: bool = False
  ) -> AsyncGenerator[LlmResponse, None]:
    yield LlmResponse(
        content=types.Content(role="model", parts=[types.Part(text="ok")])
    )


class EchoAgent(BaseAgent):
  """Plain BaseAgent (not a BaseNode) -> takes the legacy runner path."""

  async def _run_async_impl(
      self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    yield Event(
        invocation_id=ctx.invocation_id,
        author=self.name,
        content=types.Content(role="model", parts=[types.Part(text="ok")]),
    )


async def check(app_name: str, label: str, agent: BaseAgent) -> bool:
  session_service = InMemorySessionService()
  runner = Runner(
      app=App(
          name=app_name,
          root_agent=agent,
          resumability_config=ResumabilityConfig(is_resumable=True),
      ),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )
  await session_service.create_session(
      app_name=app_name, user_id=USER, session_id=SESSION
  )

  # Turn 1: a normal run, so there is an invocation to resume.
  async for _ in runner.run_async(
      user_id=USER,
      session_id=SESSION,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
  ):
    pass

  session = await session_service.get_session(
      app_name=app_name, user_id=USER, session_id=SESSION
  )
  invocation_id = session.events[0].invocation_id

  # Turn 2: resume by invocation_id, with a state_delta and NO new_message.
  async for _ in runner.run_async(
      user_id=USER,
      session_id=SESSION,
      invocation_id=invocation_id,
      state_delta={"resumed_key": "resumed_value"},
  ):
    pass

  session = await session_service.get_session(
      app_name=app_name, user_id=USER, session_id=SESSION
  )
  applied = session.state.get("resumed_key") == "resumed_value"
  print(f"  {label:<28} state={dict(session.state)!r:<34} applied={applied}")
  return applied


async def main():
  print("resume by invocation_id + state_delta, no new_message:")
  a = await check(
      "app_node", "A) node path (LlmAgent)", LlmAgent(name="root", model=EchoLlm())
  )
  b = await check(
      "app_legacy", "B) legacy path (BaseAgent)", EchoAgent(name="root")
  )
  print(f"\nRESULT: {'PASS' if (a and b) else 'FAIL (state_delta dropped)'}")


asyncio.run(main())