Skip to content

Latest commit

 

History

History
811 lines (595 loc) · 24.5 KB

File metadata and controls

811 lines (595 loc) · 24.5 KB

Developer Guide

Table of Contents


Architecture Overview

toolthread runs task handlers in background threads and lets callers (typically AI agents) poll for progress using a long-polling pattern. Here's the full flow:

Agent                          TaskManager                     Background Thread
  |                                |                                |
  |-- start_task("deploy") ------->|                                |
  |<-- task_id --------------------|-- spawns thread -------------->|
  |                                |                                |-- handler(ctx, params)
  |                                |                                |-- ctx.update("Pulling...")
  |-- poll_task(task_id) --------->|                                |
  |<-- PollResponse(updates) ------|                                |
  |                                |                                |-- ctx.ask_customer("Drop DB?")
  |-- poll_task(task_id) --------->|                                |   (handler blocks)
  |<-- PollResponse(question) -----|                                |
  |                                |                                |
  |  [Agent asks customer,         |                                |
  |   gets their answer]           |                                |
  |                                |                                |
  |-- respond_to_question() ------>|-- unblocks handler ----------->|
  |                                |                                |-- ctx.update("Migrating...")
  |-- poll_task(task_id) --------->|                                |-- return "Done"
  |<-- PollResponse(result) -------|                                |
  |                                |                                |

Key design decisions:

  • Threading, not async. Handlers run in plain threading.Threads. This keeps the API simple and compatible with any sync code your handlers need to call.
  • Long-polling, not websockets. The poll call blocks (up to a configurable timeout) until there's something new to report. This is simple, firewall-friendly, and maps perfectly to tool-call semantics.
  • Pluggable storage. The TaskStore protocol defines what toolthread needs. Swap in SQLite for persistence, or implement your own backend.

Core Concepts

TaskManager

The central orchestrator. You create one, register handlers on it, start tasks, and poll for results.

from toolthread import TaskManager

manager = TaskManager()          # uses MemoryStore by default
manager = TaskManager(store=my_store)  # or bring your own

TaskContext

The handle passed to every task handler. It provides methods to:

  • Post progress updates (ctx.update())
  • Ask the customer questions (ctx.ask_customer())
  • Check for cancellation (ctx.check_cancelled() / ctx.is_cancelled)

Handlers never interact with the store or manager directly -- TaskContext is their entire API.

TaskStore

A protocol (interface) that any storage backend must implement. Ships with two built-in implementations:

Store Persistence Dependencies Best for
MemoryStore None (in-process) None Development, testing, short-lived tasks
SQLiteStore File-based None (stdlib) Production, tasks that survive restarts

PollResponse

The object returned by poll_task(). Contains everything the caller needs:

@dataclass(frozen=True)
class PollResponse:
    task_id: str
    status: TaskStatus           # PENDING, RUNNING, WAITING_FOR_INPUT, COMPLETED, FAILED, CANCELLED
    updates: list[TaskUpdate]    # new progress messages since last poll
    pending_question: TaskQuestion | None  # question awaiting customer answer
    result: TaskResult | None    # final result (only when terminal)

Call .to_dict() on it to get a JSON-serializable dictionary.


Getting Started

Installation

pip install toolthread

For SQLite support, no extras needed -- it uses Python's built-in sqlite3. For Google ADK integration:

pip install toolthread[adk]

For development:

pip install toolthread[dev]

Basic Usage

Here's a complete working example:

import time
from toolthread import TaskManager, TaskContext

# 1. Create a manager
manager = TaskManager()

# 2. Register a handler
@manager.register_task("analyze_logs")
def analyze_logs(ctx: TaskContext, params: dict) -> dict:
    files = params.get("files", [])
    ctx.update(f"Scanning {len(files)} log files...")

    results = []
    for i, f in enumerate(files):
        ctx.check_cancelled()  # cooperative cancellation
        ctx.update(f"Processing {f} ({i+1}/{len(files)})")
        time.sleep(1)  # simulate work
        results.append({"file": f, "errors": 0})

    ctx.update("Analysis complete")
    return {"files_processed": len(files), "results": results}

# 3. Start a task
task_id = manager.start_task("analyze_logs", files=["app.log", "db.log", "web.log"])

# 4. Poll for updates
last_timestamp = None
while True:
    response = manager.poll_task(task_id, since=last_timestamp)

    for update in response.updates:
        print(f"  [{update.level.value}] {update.message}")
        last_timestamp = update.timestamp.isoformat()

    if response.result:
        print(f"\nResult: {response.result.result}")
        break

# 5. Clean up
manager.shutdown()

Writing Task Handlers

A handler is any callable with the signature (ctx: TaskContext, params: dict) -> Any. Register it with @manager.register_task("task_type").

Simple handler

@manager.register_task("ping")
def ping(ctx: TaskContext, params: dict) -> str:
    return f"pong: {params.get('message', '')}"

Handler with progress updates

Use ctx.update() to post messages that the poller will see:

@manager.register_task("data_export")
def data_export(ctx: TaskContext, params: dict) -> str:
    ctx.update("Querying database...")
    rows = fetch_all_rows()

    ctx.update(f"Exporting {len(rows)} rows to CSV...")
    write_csv(rows, params["output_path"])

    ctx.update("Uploading to S3...")
    upload_to_s3(params["output_path"], params["bucket"])

    return f"Exported {len(rows)} rows to s3://{params['bucket']}/export.csv"

Updates support severity levels:

ctx.update("Everything looks good", level="info")      # default
ctx.update("Row count seems low", level="warning")
ctx.update("Connection failed, retrying...", level="error")
ctx.update("Cache hit ratio: 0.94", level="debug")

Handler with customer interaction

Use ctx.ask_customer() to pause execution and wait for the customer's answer:

@manager.register_task("deploy")
def deploy(ctx: TaskContext, params: dict) -> str:
    ctx.update("Running pre-deploy checks...")
    issues = run_checks()

    if issues:
        ctx.update(f"Found {len(issues)} issues", level="warning")
        answer = ctx.ask_customer(
            f"Pre-deploy found {len(issues)} issues:\n"
            + "\n".join(f"- {i}" for i in issues)
            + "\n\nProceed anyway? (yes/no)",
            timeout=300.0,  # wait up to 5 minutes
        )
        if answer.strip().lower() != "yes":
            return "Deploy aborted by customer"

    ctx.update("Deploying...")
    run_deploy(params["env"])
    return f"Deployed to {params['env']}"

When ask_customer() is called:

  1. The task status changes to WAITING_FOR_INPUT
  2. The next poll_task() call returns the question in pending_question
  3. The agent shows the question to the customer
  4. The agent calls manager.respond_to_question(task_id, question_id, answer)
  5. The handler unblocks and continues with the answer

Handling cancellation

Handlers should cooperatively check for cancellation in long-running loops:

@manager.register_task("batch_process")
def batch_process(ctx: TaskContext, params: dict) -> dict:
    processed = 0
    for item in get_items():
        ctx.check_cancelled()  # raises TaskCancelled if cancelled
        process(item)
        processed += 1
        ctx.update(f"Processed {processed} items")
    return {"processed": processed}

You can also check ctx.is_cancelled without raising:

if ctx.is_cancelled:
    # do cleanup first
    cleanup()
    ctx.check_cancelled()  # now raise

The caller cancels a task with:

manager.cancel_task(task_id)

Error handling in handlers

Unhandled exceptions in handlers are caught automatically and stored as failed results:

@manager.register_task("risky_operation")
def risky_operation(ctx: TaskContext, params: dict) -> str:
    # If this raises, the task status becomes FAILED and
    # result.error contains "ValueError: something went wrong"
    raise ValueError("something went wrong")

The poller sees:

response = manager.poll_task(task_id)
assert response.status == TaskStatus.FAILED
assert response.result.error == "ValueError: something went wrong"

For expected errors you want to handle gracefully, use try/except inside your handler and return a meaningful result instead.


Storage Backends

MemoryStore

The default. Fast, zero-config, no persistence.

from toolthread import TaskManager

manager = TaskManager()  # MemoryStore is the default

Use MemoryStore when:

  • You're developing or testing
  • Tasks are short-lived and don't need to survive process restarts
  • You want the simplest possible setup

SQLiteStore

File-backed persistence using Python's built-in sqlite3. Uses WAL mode for better concurrent read/write performance.

from toolthread import TaskManager
from toolthread.store.sqlite import SQLiteStore

store = SQLiteStore("tasks.db")        # persistent file
store = SQLiteStore(":memory:")         # in-memory SQLite (for testing)

manager = TaskManager(store=store)

Use SQLiteStore when:

  • Tasks may run for a long time and you need crash recovery
  • You want to inspect task history after the fact
  • You're running in production

Custom backends

Implement the TaskStore protocol to bring your own storage (Redis, Postgres, etc.):

from toolthread.store.protocol import TaskStore

class RedisStore:
    """Your custom store -- just implement the TaskStore methods."""

    def create_task(self, task_id: str, task_type: str, params: dict) -> None:
        ...

    def get_task_status(self, task_id: str) -> TaskStatus:
        ...

    def set_task_status(self, task_id: str, status: TaskStatus) -> None:
        ...

    def add_update(self, update: TaskUpdate) -> None:
        ...

    def get_updates(self, task_id: str, since: datetime | None = None) -> list[TaskUpdate]:
        ...

    def add_question(self, question: TaskQuestion) -> None:
        ...

    def get_pending_question(self, task_id: str) -> TaskQuestion | None:
        ...

    def answer_question(self, question_id: str, answer: str) -> None:
        ...

    def set_result(self, task_id: str, result: TaskResult) -> None:
        ...

    def get_result(self, task_id: str) -> TaskResult | None:
        ...

The protocol is @runtime_checkable, so you can verify at startup:

assert isinstance(my_store, TaskStore)

All implementations must be thread-safe -- the manager writes from background threads while reads happen on the caller's thread.


Long-Polling Pattern

How it works

poll_task() doesn't return immediately when there's nothing new. Instead, it blocks for up to timeout seconds, waiting for the background thread to post an update, ask a question, or finish.

# Blocks up to 30s (default), returns immediately if there's news
response = manager.poll_task(task_id)

# Custom timeout
response = manager.poll_task(task_id, timeout=60.0)

This eliminates busy-waiting and gives you near-instant notification when something happens.

The since parameter

Pass since (an ISO-format timestamp string) to only receive updates newer than that timestamp. This prevents re-processing updates you've already seen:

last_seen = None
while True:
    response = manager.poll_task(task_id, since=last_seen)

    for update in response.updates:
        print(update.message)
        last_seen = update.timestamp.isoformat()

    if response.pending_question:
        # handle question...
        pass

    if response.result:
        break

Recommended polling pattern

def watch_task(manager, task_id):
    """Poll a task to completion, yielding updates along the way."""
    last_seen = None
    while True:
        response = manager.poll_task(task_id, since=last_seen, timeout=30.0)

        for update in response.updates:
            yield ("update", update)
            last_seen = update.timestamp.isoformat()

        if response.pending_question:
            yield ("question", response.pending_question)

        if response.result:
            yield ("result", response.result)
            return

Google ADK Integration

toolthread ships with optional Google ADK support that wraps TaskManager operations as ADK-compatible tool definitions.

Setup

pip install toolthread[adk]
from toolthread import TaskManager
from toolthread.adk import create_adk_tools

manager = TaskManager()

# Register your handlers
@manager.register_task("run_analysis")
def run_analysis(ctx, params):
    ctx.update("Starting analysis...")
    # ... work ...
    return {"findings": [...]}

# Create ADK tools
tools = create_adk_tools(manager)

Tool definitions

create_adk_tools() returns tools that map to the core TaskManager operations:

Tool Maps to Purpose
start_task manager.start_task() Start a background task
poll_task manager.poll_task() Long-poll for updates
respond_to_question manager.respond_to_question() Answer a pending question
cancel_task manager.cancel_task() Cancel a running task

Complete agent example

from toolthread import TaskManager, TaskContext
from toolthread.adk import create_adk_tools
from google.adk import Agent

# Set up task manager with handlers
manager = TaskManager()

@manager.register_task("deploy_service")
def deploy_service(ctx: TaskContext, params: dict) -> str:
    ctx.update(f"Deploying {params['service']} to {params['env']}...")
    # ... deployment logic ...
    answer = ctx.ask_customer("Deployment ready. Run smoke tests? (yes/no)")
    if answer.lower() == "yes":
        ctx.update("Running smoke tests...")
    return "Deployment complete"

# Create the agent with toolthread tools
agent = Agent(
    model="gemini-2.0-flash",
    name="deploy_agent",
    instruction=SYSTEM_INSTRUCTIONS,  # see template below
    tools=create_adk_tools(manager),
)

System Instructions Template

Copy and adapt this template for any agent that uses toolthread tools. It tells the LLM how to use the background task workflow correctly.

You have access to tools for running long-running tasks in the background. Here is how they work:

## Available Tools

- **start_task(task_type, ...)**: Start a background task. Returns a task_id.
  Use this for operations that may take more than a few seconds.

- **poll_task(task_id, since?, timeout?)**: Check on a running task.
  Returns status, new progress updates, any pending customer questions, and the
  final result if complete. This long-polls -- it will wait up to `timeout`
  seconds for new activity before returning.

- **respond_to_question(task_id, question_id, answer)**: Answer a question
  that a background task is asking the customer. The task is paused until
  you provide the answer.

- **cancel_task(task_id)**: Cancel a running task.

## Workflow

1. When asked to perform a long-running operation, call `start_task` with the
   appropriate task_type and parameters. Tell the customer you've started the
   task.

2. Immediately call `poll_task` with the returned task_id to wait for updates.

3. When poll_task returns:
   - **Updates**: Summarize new progress updates for the customer.
     Don't repeat updates they've already seen.
   - **Pending question**: The task needs customer input. Present the question
     clearly to the customer and wait for their response. Then call
     `respond_to_question` with their answer.
   - **Result**: The task is done. Present the final result to the customer.

4. After presenting updates (if the task isn't done), call `poll_task` again
   with the `since` parameter set to the timestamp of the last update you
   received. This ensures you only get new updates.

5. Repeat until you get a final result.

## Important Notes

- Always pass `since` after your first poll to avoid seeing duplicate updates.
- When a question is pending, do NOT keep polling. Present the question to the
  customer first, get their answer, then call respond_to_question, then resume
  polling.
- If a task fails, the result will contain an error message. Present it to the
  customer and suggest next steps.
- Use cancel_task if the customer wants to stop a running task.
- Don't start multiple instances of the same task unless the customer asks.

## Available Task Types

[List your registered task types here, e.g.]
- `deploy`: Deploy a service. Params: service (str), env (str)
- `analyze_logs`: Analyze log files. Params: files (list[str])

Best Practices

Polling intervals

  • The default timeout=30.0 for poll_task() is a good starting point. The long-poll returns immediately when there's news, so the timeout only matters for idle periods.
  • For ADK/tool-call agents, 30 seconds works well -- most frameworks handle this gracefully.
  • For manual integrations, you can increase the timeout to reduce round-trips, or decrease it if you need faster cancellation response.

Timeout configuration

  • ask_customer() defaults to 300 seconds (5 minutes). Adjust based on how quickly you expect customers to respond.
  • If a question times out, TaskTimeout is raised in the handler. Decide whether to fail, use a default, or skip that step.
from toolthread import TaskTimeout

@manager.register_task("setup")
def setup(ctx: TaskContext, params: dict) -> str:
    try:
        answer = ctx.ask_customer("Which region?", timeout=60.0)
    except TaskTimeout:
        answer = "us-east-1"  # sensible default
        ctx.update(f"No response, defaulting to {answer}", level="warning")
    return f"Set up in {answer}"

Graceful degradation

  • Always call manager.shutdown() when your application exits. This cancels running tasks and waits for threads to finish.
  • Handlers should check ctx.check_cancelled() or ctx.is_cancelled periodically in long loops.
  • Wrap long-running work in handlers rather than doing it inline -- this way the agent can always report progress.

Testing your handlers

Test handlers directly by creating a TaskManager with the default MemoryStore:

def test_deploy_handler():
    manager = TaskManager()

    @manager.register_task("deploy")
    def deploy(ctx: TaskContext, params: dict) -> str:
        ctx.update("Deploying...")
        return "done"

    task_id = manager.start_task("deploy", env="staging")

    # Poll until complete
    response = manager.poll_task(task_id, timeout=5.0)
    assert response.status == TaskStatus.COMPLETED
    assert response.result.result == "done"

For testing customer interaction, use a separate thread to provide answers:

import threading

def test_ask_customer():
    manager = TaskManager()

    @manager.register_task("interactive")
    def interactive(ctx: TaskContext, params: dict) -> str:
        answer = ctx.ask_customer("Continue?")
        return f"Customer said: {answer}"

    task_id = manager.start_task("interactive")

    # Poll to get the question
    response = manager.poll_task(task_id, timeout=5.0)
    assert response.pending_question is not None

    # Answer it
    manager.respond_to_question(
        task_id,
        response.pending_question.question_id,
        "yes",
    )

    # Poll for result
    response = manager.poll_task(task_id, timeout=5.0)
    assert response.result.result == "Customer said: yes"

API Reference

TaskManager

class TaskManager:
    def __init__(self, store: TaskStore | None = None) -> None

Methods:

Method Signature Description
register_task (task_type: str) -> Callable Decorator to register a handler for task_type
start_task (task_type: str, **params) -> str Start a background task, returns task_id
poll_task (task_id: str, since: str | None = None, timeout: float = 30.0) -> PollResponse Long-poll for updates
respond_to_question (task_id: str, question_id: str, answer: str) -> bool Answer a pending question
cancel_task (task_id: str) -> bool Request cancellation of a task
shutdown (timeout: float = 10.0) -> None Cancel all tasks and wait for threads

TaskContext

class TaskContext:
    task_id: str                  # read-only property
    is_cancelled: bool            # read-only property

Methods:

Method Signature Description
update (message: str, level: str | UpdateLevel = "info") -> None Post a progress update
ask_customer (prompt: str, timeout: float = 300.0) -> str Ask customer a question, blocks until answered
check_cancelled () -> None Raise TaskCancelled if task was cancelled

TaskStore protocol

@runtime_checkable
class TaskStore(Protocol):
    def create_task(self, task_id: str, task_type: str, params: dict[str, Any]) -> None: ...
    def get_task_status(self, task_id: str) -> TaskStatus: ...
    def set_task_status(self, task_id: str, status: TaskStatus) -> None: ...
    def add_update(self, update: TaskUpdate) -> None: ...
    def get_updates(self, task_id: str, since: datetime | None = None) -> list[TaskUpdate]: ...
    def add_question(self, question: TaskQuestion) -> None: ...
    def get_pending_question(self, task_id: str) -> TaskQuestion | None: ...
    def answer_question(self, question_id: str, answer: str) -> None: ...
    def set_result(self, task_id: str, result: TaskResult) -> None: ...
    def get_result(self, task_id: str) -> TaskResult | None: ...

Data types

TaskStatus

class TaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    WAITING_FOR_INPUT = "waiting_for_input"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"

    @property
    def is_terminal(self) -> bool: ...  # True for COMPLETED, FAILED, CANCELLED

TaskUpdate

@dataclass(frozen=True)
class TaskUpdate:
    task_id: str
    message: str
    level: UpdateLevel = UpdateLevel.INFO
    timestamp: datetime = ...     # auto-set to UTC now
    update_id: str = ...          # auto-generated

UpdateLevel

class UpdateLevel(Enum):
    DEBUG = "debug"
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"

TaskQuestion

@dataclass
class TaskQuestion:
    question_id: str              # auto-generated
    task_id: str
    prompt: str
    timeout: float = 300.0
    asked_at: datetime = ...      # auto-set to UTC now
    answered_at: datetime | None = None
    answer: str | None = None
    timed_out: bool = False

    @property
    def is_answered(self) -> bool: ...
    @property
    def is_resolved(self) -> bool: ...  # answered or timed out

TaskResult

@dataclass(frozen=True)
class TaskResult:
    task_id: str
    status: TaskStatus
    result: Any = None
    error: str | None = None
    completed_at: datetime = ...  # auto-set to UTC now

PollResponse

@dataclass(frozen=True)
class PollResponse:
    task_id: str
    status: TaskStatus
    updates: list[TaskUpdate] = []
    pending_question: TaskQuestion | None = None
    result: TaskResult | None = None

    def to_dict(self) -> dict[str, Any]: ...  # JSON-serializable output

Exceptions

Exception Raised when
TaskCancelled Handler checks cancellation and the task was cancelled
TaskTimeout A ctx.ask_customer() call times out without an answer
TaskNotFound A task_id or question_id doesn't exist in the store