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

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

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

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