Session Management Examples

This directory contains examples demonstrating the complete lifecycle of AGB sessions.

Key Concepts

Session Creation

• Use CreateSessionParams to configure your session (image, timeouts).
• Always check result.success before proceeding.

Session Retrieval

• You can reconnect to any active session using agb.get(session_id).
• This is useful for stateless applications (e.g., web servers) that need to resume control of a session.

Session Pooling

• For high-throughput applications, creating a new session for every request is inefficient.
• Maintaining a pool of "warm" sessions allows for faster execution and better resource utilization.

Cleanup

• Always call agb.delete(session) in a finally block or when the session is no longer needed to avoid unnecessary charges.

Examples

Basic Session Lifecycle (create_session.py)

Basic session creation, listing, and deletion.

"""
This example demonstrates how to create a session in AGB.
"""

import os
import time
from typing import Dict, Optional

from agb import AGB
from agb.session_params import CreateSessionParams


def create_session_with_default_params() -> None:
    """Create a session with default parameters."""
    # Initialize the AGB client
    api_key = os.environ.get("AGB_API_KEY", "")
    agb = AGB(api_key=api_key)

    # Create a session with default parameters
    params = CreateSessionParams(image_id="agb-code-space-1")
    result = agb.create(params)

    if result.success and result.session:
        session = result.session
        print(f"Session created successfully with ID: {session.session_id}")
        print(f"Request ID: {result.request_id}")

        # Clean up
        delete_result = agb.delete(session)
        if delete_result.success:
            print("Session deleted successfully")
        else:
            print(f"Failed to delete session: {delete_result.error_message}")
    else:
        print(f"Failed to create session: {result.error_message}")


def create_session_with_custom_image() -> None:
    """Create a session with custom image."""
    # Initialize the AGB client
    api_key = os.environ.get("AGB_API_KEY", "")
    agb = AGB(api_key=api_key)

    # Create session parameters with custom image
    params = CreateSessionParams(image_id="agb-code-space-1")

    # Create a session with the parameters
    result = agb.create(params)

    if result.success and result.session:
        session = result.session
        print(f"Session with custom image created successfully with ID: {session.session_id}")
        print(f"Request ID: {result.request_id}")

        # Clean up
        delete_result = agb.delete(session)
        if delete_result.success:
            print("Session deleted successfully")
        else:
            print(f"Failed to delete session: {delete_result.error_message}")
    else:
        print(f"Failed to create session with custom image: {result.error_message}")

def create_session_with_labels() -> None:
    """Create a session with labels."""
    # Initialize the AGB client
    api_key = os.environ.get("AGB_API_KEY", "")
    agb = AGB(api_key=api_key)

    # Define labels
    labels: Dict[str, str] = {
        "environment": "development",
        "project": "example",
        "owner": "user123",
        "team": "backend",
        "version": "v1.0.0"
    }

    # Create session parameters with labels
    params = CreateSessionParams(
        image_id="agb-browser-use-1",
        labels=labels
    )

    # Create a session with the parameters
    result = agb.create(params)

    if result.success and result.session:
        session = result.session
        print(f"Session with labels created successfully with ID: {session.session_id}")
        print(f"Request ID: {result.request_id}")

        # Verify the labels were set
        label_result = session.get_labels()
        if label_result.success:
            retrieved_labels = label_result.data
            print("Retrieved labels:")
            for key, value in retrieved_labels.items():
                print(f"  {key}: {value}")
        else:
            print(f"Failed to get labels: {label_result.error_message}")

        # Update labels during session lifecycle
        updated_labels = {
            **labels,
            "status": "active",
            "last_updated": str(int(time.time()))
        }

        update_result = session.set_labels(updated_labels)
        if update_result.success:
            print("Labels updated successfully")

            # Get updated labels
            final_label_result = session.get_labels()
            if final_label_result.success:
                final_labels = final_label_result.data
                print("Final labels:")
                for key, value in final_labels.items():
                    print(f"  {key}: {value}")
        else:
            print(f"Failed to update labels: {update_result.error_message}")

        # Clean up
        delete_result = agb.delete(session)
        if delete_result.success:
            print("Session deleted successfully")
        else:
            print(f"Failed to delete session: {delete_result.error_message}")
    else:
        print(f"Failed to create session with labels: {result.error_message}")


def main() -> None:
    """Run all examples."""
    print("1. Creating session with default parameters...")
    create_session_with_default_params()
    print("\n2. Creating session with custom image...")
    create_session_with_custom_image()
    print("\n3. Creating session with labels...")
    create_session_with_labels()


if __name__ == "__main__":
    main()

Reconnecting to Sessions (get_session.py)

How to retrieve an existing session by ID.

"""
Example demonstrating how to use the Get API to retrieve a session by its ID.

This example shows:
1. Creating a session
2. Retrieving the session using the Get API
3. Using the session for operations
4. Cleaning up resources
"""

import os
import sys

# Add the project root to Python path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..'))
sys.path.insert(0, project_root)

from agb import AGB
from agb.session_params import CreateSessionParams

def main():
    # Get API key from environment variable
    api_key = os.getenv("AGB_API_KEY")
    if not api_key:
        raise ValueError("AGB_API_KEY environment variable is not set")

    # Initialize AGB client
    agb = AGB(api_key=api_key)

    # For demonstration, first create a session
    print("Creating a session...")
    params = CreateSessionParams(image_id="agb-browser-use-1")
    create_result = agb.create(params)
    if not create_result.success:
        raise RuntimeError(f"Failed to create session: {create_result.error_message}")

    session_id = create_result.session.session_id
    print(f"Created session with ID: {session_id}")

    # Retrieve the session by ID using Get API
    print("\nRetrieving session using Get API...")
    get_result = agb.get(session_id)

    if not get_result.success:
        raise RuntimeError(f"Failed to get session: {get_result.error_message}")

    session = get_result.session

    # Display session information
    print("Successfully retrieved session:")
    print(f"  Session ID: {session.session_id}")
    print(f"  Request ID: {get_result.request_id}")
    print(f"  Resource URL: {session.resource_url}")
    print(f"  Resource ID: {session.resource_id}")
    print(f"  App Instance ID: {session.app_instance_id}")

    # The session object can be used for further operations
    # For example, you can execute commands, work with files, etc.
    print("\nSession is ready for use")

    # Clean up: Delete the session when done
    print("\nCleaning up...")
    delete_result = agb.delete(session)

    if delete_result.success:
        print(f"Session {session_id} deleted successfully")
    else:
        print(f"Failed to delete session {session_id}")

if __name__ == "__main__":
    main()

Session Pooling (session_pool.py)

Implementation of a thread-safe session pool for high-concurrency applications.

"""
AGB Session Pool Example

This example demonstrates how to implement a thread-safe session pool.
Pooling sessions is critical for applications that need to execute many short-lived operations
without incurring the latency of creating a new session for every request.
"""

import os
import threading
import time
from contextlib import contextmanager
from typing import Dict, Optional

from agb import AGB
from agb.session_params import CreateSessionParams

class SessionPool:
    """Thread-safe session pool for high-throughput applications"""

    def __init__(self, api_key: str, max_sessions: int = 5, session_timeout: int = 300):
        self.agb = AGB(api_key=api_key)
        self.max_sessions = max_sessions
        self.session_timeout = session_timeout
        self.sessions: Dict[str, dict] = {}
        self.lock = threading.Lock()
        print(f"🎱 Session Pool initialized (max={max_sessions})")

    @contextmanager
    def get_session(self):
        """Context manager to acquire and release a session"""
        session_info = self._acquire_session()
        try:
            print(f"🟢 Acquired session: {session_info['id']}")
            yield session_info['session']
        finally:
            print(f"🟡 Released session: {session_info['id']}")
            self._release_session(session_info['id'])

    def _acquire_session(self):
        with self.lock:
            # 1. Cleanup expired sessions
            self._cleanup_expired_sessions()

            # 2. Try to reuse an idle session
            for session_id, info in self.sessions.items():
                if not info['in_use']:
                    info['in_use'] = True
                    info['last_used'] = time.time()
                    return info

            # 3. Create new session if under limit
            if len(self.sessions) < self.max_sessions:
                print("✨ Creating new session for pool...")
                params = CreateSessionParams(image_id="agb-code-space-1")
                result = self.agb.create(params)

                if result.success:
                    session_info = {
                        'id': result.session.session_id,
                        'session': result.session,
                        'created': time.time(),
                        'last_used': time.time(),
                        'in_use': True
                    }
                    self.sessions[session_info['id']] = session_info
                    return session_info
                else:
                    raise Exception(f"Failed to create session: {result.error_message}")

            raise Exception("No sessions available and pool is at maximum capacity")

    def _release_session(self, session_id: str):
        with self.lock:
            if session_id in self.sessions:
                self.sessions[session_id]['in_use'] = False

    def _cleanup_expired_sessions(self):
        current_time = time.time()
        expired_ids = []

        for session_id, info in self.sessions.items():
            # If idle for too long
            if not info['in_use'] and (current_time - info['last_used']) > self.session_timeout:
                expired_ids.append(session_id)

        for session_id in expired_ids:
            print(f"⌛ Cleaning up expired session: {session_id}")
            session_info = self.sessions.pop(session_id)
            try:
                self.agb.delete(session_info['session'])
            except Exception as e:
                print(f"Error deleting session {session_id}: {e}")

    def destroy_all(self):
        """Clean up all sessions in the pool"""
        print("💥 Destroying pool...")
        with self.lock:
            for session_info in self.sessions.values():
                try:
                    self.agb.delete(session_info['session'])
                    print(f"Deleted {session_info['id']}")
                except Exception as e:
                    print(f"Error deleting session: {e}")
            self.sessions.clear()

def main():
    api_key = os.getenv("AGB_API_KEY")
    if not api_key:
        print("Error: AGB_API_KEY environment variable is not set")
        return

    # 1. Initialize pool
    pool = SessionPool(api_key=api_key, max_sessions=2)

    try:
        # 2. Run tasks using the pool
        # We'll simulate 4 tasks with a pool of size 2
        # This means sessions will be reused

        def run_task(task_name):
            try:
                with pool.get_session() as session:
                    print(f"▶️  Running {task_name}...")
                    result = session.code.run_code(f"print('Hello from {task_name}')", "python")
                    # Simulate work holding the session
                    time.sleep(1)
                    print(f"✅ {task_name} Result: {result.result.strip()}")
            except Exception as e:
                print(f"❌ {task_name} failed: {e}")

        threads = []
        for i in range(4):
            t = threading.Thread(target=run_task, args=(f"Task-{i}",))
            threads.append(t)
            t.start()

        for t in threads:
            t.join()

    finally:
        # 3. Cleanup
        pool.destroy_all()

if __name__ == "__main__":
    main()

How to Run

1. Set your API Key:

   export AGB_API_KEY="your_api_key_here"

2. Run Examples:

   python docs/examples/session_management/create_session.py
   python docs/examples/session_management/get_session.py
   python docs/examples/session_management/session_pool.py