Demonstrates session rewind.
()
| 88 | |
| 89 | |
| 90 | async def main(): |
| 91 | """Demonstrates session rewind.""" |
| 92 | print("🚀 Testing Rewind Session Feature") |
| 93 | print("=" * 50) |
| 94 | |
| 95 | runner = InMemoryRunner( |
| 96 | agent=agent.root_agent, |
| 97 | app_name=APP_NAME, |
| 98 | ) |
| 99 | |
| 100 | # Create a session |
| 101 | session = await runner.session_service.create_session( |
| 102 | app_name=APP_NAME, user_id=USER_ID |
| 103 | ) |
| 104 | print(f"Created session: {session.id}") |
| 105 | |
| 106 | # 1. Initial agent calls to set state and artifact |
| 107 | print("\n\n===== INITIALIZING STATE AND ARTIFACT =====") |
| 108 | await call_agent_async( |
| 109 | runner, USER_ID, session.id, "set state `color` to red" |
| 110 | ) |
| 111 | await call_agent_async( |
| 112 | runner, USER_ID, session.id, "save artifact file1 with content version1" |
| 113 | ) |
| 114 | |
| 115 | # 2. Check current state and artifact |
| 116 | print("\n\n===== STATE BEFORE UPDATE =====") |
| 117 | await call_agent_async( |
| 118 | runner, USER_ID, session.id, "what is the value of state `color`?" |
| 119 | ) |
| 120 | await call_agent_async(runner, USER_ID, session.id, "load artifact file1") |
| 121 | |
| 122 | # 3. Update state and artifact - THIS IS THE POINT WE WILL REWIND BEFORE |
| 123 | print("\n\n===== UPDATING STATE AND ARTIFACT =====") |
| 124 | events_update_state = await call_agent_async( |
| 125 | runner, USER_ID, session.id, "update state key color to blue" |
| 126 | ) |
| 127 | rewind_invocation_id = events_update_state[0].invocation_id |
| 128 | print(f"Will rewind before invocation: {rewind_invocation_id}") |
| 129 | |
| 130 | await call_agent_async( |
| 131 | runner, USER_ID, session.id, "save artifact file1 with content version2" |
| 132 | ) |
| 133 | |
| 134 | # 4. Check state and artifact after update |
| 135 | print("\n\n===== STATE AFTER UPDATE =====") |
| 136 | await call_agent_async( |
| 137 | runner, USER_ID, session.id, "what is the value of state key color?" |
| 138 | ) |
| 139 | await call_agent_async(runner, USER_ID, session.id, "load artifact file1") |
| 140 | |
| 141 | # 5. Perform rewind |
| 142 | print(f"\n\n===== REWINDING SESSION to before {rewind_invocation_id} =====") |
| 143 | await runner.rewind_async( |
| 144 | user_id=USER_ID, |
| 145 | session_id=session.id, |
| 146 | rewind_before_invocation_id=rewind_invocation_id, |
| 147 | ) |
no test coverage detected