MODEL = "gemini-2.0-flash-001"
AGENT_APP_NAME = 'bqassistant'

GOOGLE_CLIENT_ID = os.getenv('GOOGLE_CLIENT_ID')
GOOGLE_CLIENT_SECRET = os.getenv('GOOGLE_CLIENT_SECRET')

bqcredentials = BigQueryCredentialsConfig(
        client_id = GOOGLE_CLIENT_ID, 
        client_secret = GOOGLE_CLIENT_SECRET
)

bqtoolset = BigQueryToolset(
        credentials_config=bqcredentials
)

bqassistant = Agent(
        model=MODEL,
        name="bqassistant",
        description="Use BigQUery tools to handle user queries",
        instruction="You are helpful BigQuery assistant",
        tools = [bqtoolset],
)

bq_recommender = Agent(
        model=MODEL,
        name="bqrecommender",
        description="Generate next 10 recommendations on what you can do with BigQuery",
        instruction="Generate next 10 recommendations on what you can do with BigQuery"
)


root_agent = SequentialAgent(
    name = "sequence",
    sub_agents = [bqassistant, bq_recommender]
)


is_final response flag on event does not take into account the state of oauth process. 

As workaround I created custom agent - as according to documentation this is one of the workarounds to code custom flows assuming I would be able to block transition to next agent in sequence untill state variable with bigquery_access_token is set. 

however, this way second agent is never executed. simply when first agent starts oauth2 flow and flags event with is final response = true ---> it passes flow to code that checks if state variable is set and if it is not set it simply skipps next steps - in other words -- it does not stop the flow .. it just skipps it. 


So i modified custom agent to add pooling - wait in loop util state key is set or it timesout (lets say 6 minutes for user to go trhough oauth flow). Seems it blocks the whole agent from even starting oauth flow:

      @override
      async def _run_async_impl(
                self, ctx: InvocationContext
      ) -> AsyncGenerator[Event, None]:

        logger.info(f"[{self.name}] Starting story generation workflow.")
        async for event in self.bigquery_oauth_task.run_async(ctx):
            logger.info(f"[{self.name}] Event from First Agent: {event.model_dump_json(indent=2, exclude_none=True)}")
            yield event
        
        logger.info(f"[{self.name}] Waiting for OAuth token with a {TIMEOUT_SECONDS}s timeout...")
        start_time = time.monotonic()
        token_found = False
        while time.monotonic() - start_time < TIMEOUT_SECONDS:
                # 1. Check if the token has been successfully set in the session state
                if "bigquery_access_token" in ctx.session.state and ctx.session.state["bigquery_access_token"]:
                        logger.info(f"[{self.name}] BigQuery access token successfully acquired.")
                        token_found = True
                        break # Exit the polling loop


                # 3. Wait asynchronously before the next check
                await asyncio.sleep(POLLING_INTERVAL_SECONDS)
                
        # After the loop, verify that the token was actually found before proceeding
        if not token_found:
                logger.error(f"[{self.name}] Timed out waiting for BigQuery access token. Aborting workflow.")
                return # Stop processing if we timed out

        # 2. next step
        logger.info(f"[{self.name}] Next agent")
        # Use the loop_agent instance attribute assigned during init
        async for event in self.generic_consultaing_task.run_async(ctx):
            logger.info(f"[{self.name}] Event from Second Agent: {event.model_dump_json(indent=2, exclude_none=True)}")
            yield event

        logger.info(f"[{self.name}] Workflow finished.")