feat: Client updates for Axon SSE methods

AGENTS.md

@@ -0,0 +1,154 @@
+# Agent Workflow Guidelines
+
+## Making Code Changes
+
+When making changes to the codebase, follow this workflow to ensure code quality:
+
+### 1. Run Lint Before Changes
+
+Before making any code changes, run the linter to establish a baseline:
+
+```bash
+./scripts/lint
+```
+
+This runs:
+- `pyright` - Type checking
+- `mypy` - Additional type checking
+- `ruff check` - Code linting
+- `ruff format --check` - Format checking
+- Import validation
+
+### 2. Make Your Changes
+
+Make the necessary code changes, ensuring you:
+- Follow existing code patterns
+- Update type annotations
+- Handle edge cases (e.g., `Omit`, `NotGiven`, `None`)
+- Maintain backward compatibility
+
+### 3. Update Tests
+
+Update or add tests for your changes:
+- Unit tests in `tests/`
+- Smoke tests in `tests/smoketests/`
+- Ensure tests match the new behavior
+
+### 4. Run Lint After Changes
+
+After making changes, run the linter again to catch any issues:
+
+```bash
+./scripts/lint
+```
+
+Fix any new errors or warnings that appear.
+
+### 5. Run Format
+
+Apply code formatting to ensure consistent style:
+
+```bash
+ruff format .
+```
+
+This auto-formats all Python files to match the project's style guidelines.
+
+### 6. Run Tests
+
+Run the test suite to ensure everything works:
+
+```bash
+# Run all tests
+uv run pytest
+
+# Run specific tests
+uv run pytest tests/test_axon_sse_reconnect.py -xvs
+
+# Run smoke tests
+uv run pytest tests/smoketests/ -m smoketest
+```
+
+### 7. Commit Changes
+
+Once lint, format, and tests pass, commit your changes:
+
+```bash
+git add -A
+git commit -m "type: description
+
+Detailed explanation of changes
+
+Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>"
+```
+
+## Common Patterns
+
+### Handling Optional Parameters
+
+When dealing with parameters that can be omitted:
+
+```python
+# Parameter definition
+def method(self, param: int | Omit = omit):
+    ...
+
+# In implementation, check for Omit before using
+if not isinstance(param, Omit):
+    use_param(param)
+else:
+    # param was omitted, use default behavior
+    pass
+```
+
+### Type-Safe Transformations
+
+The `transform` function automatically filters out `Omit` and `NotGiven` values:
+
+```python
+# This dict with omitted values
+{"field": omit, "other": 123}
+
+# Becomes this after transform
+{"other": 123}
+```
+
+### Testing with Mocks
+
+When testing methods that call `_get` or similar:
+
+```python
+with patch.object(client.resource, "_get") as mock_get:
+    mock_stream = Mock(spec=Stream)
+    mock_get.return_value = mock_stream
+
+    # Your test code
+    result = client.resource.method()
+
+    # Verify the call
+    call_args = mock_get.call_args
+    options = call_args.kwargs["options"]
+    assert options["params"]["field"] == expected_value
+```
+
+## Merge Conflict Resolution
+
+When merging branches:
+
+1. **Understand both changes** - Read the diff from both sides
+2. **Combine features intelligently** - Don't just pick one side
+3. **Update tests** - Tests may need adjustments for merged behavior
+4. **Fix type errors** - Merges can introduce type mismatches
+5. **Validate syntax** - Run `python3 -m py_compile` on changed files
+6. **Run full lint** - Ensure everything passes
+
+## Example: Recent Merge Fix
+
+The merge of `origin/main` into `feature/ts-pr-765-port` required:
+
+1. **Code fix** - Handle `Omit` type properly in reconnection logic
+2. **Test fix** - Update assertions from `is None` to `not in dict`
+3. **Syntax check** - Verify Python syntax is valid
+4. **Type check** - Ensure mypy/pyright pass
+
+See commit history for the resolution pattern.

scripts/generate_examples_md.py

@@ -11,10 +11,10 @@
 import re
 import sys
 import argparse
-from typing import Any
+from typing import Any, cast
 from pathlib import Path
 
-import frontmatter  # type: ignore[import-untyped]
+import frontmatter  # type: ignore[import-not-found, import-untyped]
 
 ROOT = Path(__file__).parent.parent
 EXAMPLES_DIR = ROOT / "examples"
@@ -38,7 +38,7 @@ def parse_example(path: Path) -> dict[str, Any]:
         raise ValueError(f"{path}: docstring must start with frontmatter (---)")
 
     try:
-        post = frontmatter.loads(docstring)
+        post = cast(Any, frontmatter).loads(docstring)
         return dict(post.metadata)
     except Exception as e:
         raise ValueError(f"{path}: invalid frontmatter: {e}") from e

src/runloop_api_client/_utils/_compat.py

@@ -2,7 +2,7 @@
 
 import sys
 import typing_extensions
-from typing import Any, Type, Union, Literal, Optional
+from typing import Any, Type, Union, Literal, Optional, cast
 from datetime import date, datetime
 from typing_extensions import get_args as _get_args, get_origin as _get_origin
 
@@ -34,7 +34,8 @@ def is_typeddict(tp: Type[Any]) -> bool:
 
 
 def is_literal_type(tp: Type[Any]) -> bool:
-    return get_origin(tp) in _LITERAL_TYPES
+    origin = get_origin(tp)
+    return cast(Any, origin) in _LITERAL_TYPES
 
 
 def parse_date(value: Union[date, StrBytesIntFloat]) -> date:

src/runloop_api_client/_utils/_utils.py

@@ -372,8 +372,7 @@ def file_from_path(path: str) -> FileTypes:
 def get_required_header(headers: HeadersLike, header: str) -> str:
     lower_header = header.lower()
     if is_mapping_t(headers):
-        # mypy doesn't understand the type narrowing here
-        for k, v in headers.items():  # type: ignore
+        for k, v in cast(Mapping[str, object], headers).items():
             if k.lower() == lower_header and isinstance(v, str):
                 return v
 

src/runloop_api_client/lib/cancellation.py

@@ -0,0 +1,103 @@
+"""Cancellation support for polling operations."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from typing import TYPE_CHECKING
+
+from .._exceptions import RunloopError
+
+if TYPE_CHECKING:
+    pass
+
+__all__ = ["PollingCancelled", "CancellationToken"]
+
+
+class PollingCancelled(RunloopError):
+    """Exception raised when a polling operation is cancelled."""
+
+    pass
+
+
+class CancellationToken:
+    """Thread-safe cancellation token for polling operations.
+
+    Similar to JavaScript's AbortSignal. Works in both sync and async contexts.
+
+    Example (sync):
+        >>> token = CancellationToken()
+        >>> # In another thread:
+        >>> token.cancel()
+        >>> # In polling code:
+        >>> token.raise_if_cancelled()  # Raises PollingCancelled
+
+    Example (async):
+        >>> token = CancellationToken()
+        >>> # In another task:
+        >>> token.cancel()
+        >>> # In async polling code:
+        >>> await asyncio.wait_for(token.async_event.wait(), timeout=1.0)
+    """
+
+    def __init__(self) -> None:
+        """Create a new cancellation token."""
+        self._cancelled = False
+        self._sync_event = threading.Event()
+        self._async_event: asyncio.Event | None = None
+        self._lock = threading.Lock()
+
+    def cancel(self) -> None:
+        """Mark this token as cancelled.
+
+        Thread-safe and can be called multiple times. Sets both sync and async events.
+        """
+        with self._lock:
+            if self._cancelled:
+                return
+            self._cancelled = True
+            self._sync_event.set()
+            if self._async_event is not None:
+                self._async_event.set()
+
+    def is_cancelled(self) -> bool:
+        """Check if this token has been cancelled.
+
+        Returns:
+            True if cancel() has been called, False otherwise.
+        """
+        return self._cancelled
+
+    def raise_if_cancelled(self) -> None:
+        """Raise PollingCancelled if this token has been cancelled.
+
+        Raises:
+            PollingCancelled: If cancel() has been called.
+        """
+        if self._cancelled:
+            raise PollingCancelled("Polling operation was cancelled")
+
+    @property
+    def sync_event(self) -> threading.Event:
+        """Get the synchronous event for cancellation checking.
+
+        Returns:
+            threading.Event that is set when cancel() is called.
+        """
+        return self._sync_event
+
+    @property
+    def async_event(self) -> asyncio.Event:
+        """Get the asynchronous event for cancellation checking.
+
+        Lazily creates the async event on first access. If cancel() was already called,
+        the event will be set immediately.
+
+        Returns:
+            asyncio.Event that is set when cancel() is called.
+        """
+        if self._async_event is None:
+            self._async_event = asyncio.Event()
+            if self._cancelled:
+                self._async_event.set()
+        return self._async_event

src/runloop_api_client/lib/polling.py

@@ -2,6 +2,8 @@
 from typing import Any, TypeVar, Callable, Optional
 from dataclasses import dataclass
 
+from .cancellation import CancellationToken
+
 T = TypeVar("T")
 
 
@@ -27,6 +29,7 @@ def poll_until(
     is_terminal: Callable[[T], bool],
     config: Optional[PollingConfig] = None,
     on_error: Optional[Callable[[Exception], T]] = None,
+    cancellation_token: Optional[CancellationToken] = None,
 ) -> T:
     """
     Poll until a condition is met or timeout/max attempts are reached.
@@ -37,12 +40,14 @@ def poll_until(
         config: Optional polling configuration
         on_error: Optional error handler that can return a value to continue polling
                  or re-raise the exception to stop polling
+        cancellation_token: Optional token to cancel the polling operation
 
     Returns:
         The final state of the polled object
 
     Raises:
         PollingTimeout: When max attempts or timeout is reached
+        PollingCancelled: If cancellation_token.cancel() is called
     """
     if config is None:
         config = PollingConfig()
@@ -52,6 +57,10 @@ def poll_until(
     last_result = None
 
     while True:
+        # Check for cancellation before each iteration
+        if cancellation_token is not None:
+            cancellation_token.raise_if_cancelled()
+
         try:
             last_result = retriever()
         except Exception as e:
@@ -72,4 +81,9 @@ def poll_until(
             if elapsed >= config.timeout_seconds:
                 raise PollingTimeout(f"Exceeded timeout of {config.timeout_seconds} seconds", last_result)
 
-        time.sleep(config.interval_seconds)
+        # Cancellable sleep
+        if cancellation_token is not None:
+            if cancellation_token.sync_event.wait(timeout=config.interval_seconds):
+                cancellation_token.raise_if_cancelled()
+        else:
+            time.sleep(config.interval_seconds)