agentabi 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {agentabi-0.2.0/src/agentabi.egg-info → agentabi-0.2.2}/PKG-INFO +10 -5
  2. {agentabi-0.2.0 → agentabi-0.2.2}/README.md +7 -2
  3. {agentabi-0.2.0 → agentabi-0.2.2}/pyproject.toml +2 -2
  4. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/__init__.py +24 -1
  5. agentabi-0.2.2/src/agentabi/middleware/__init__.py +119 -0
  6. agentabi-0.2.2/src/agentabi/middleware/logging.py +112 -0
  7. agentabi-0.2.2/src/agentabi/middleware/timeout.py +84 -0
  8. agentabi-0.2.2/src/agentabi/middleware/usage.py +134 -0
  9. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/base.py +51 -1
  10. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/claude_native.py +55 -1
  11. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/codex_native.py +72 -18
  12. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/codex_sdk.py +45 -25
  13. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/gemini_native.py +48 -7
  14. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/opencode_native.py +51 -13
  15. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/session.py +94 -3
  16. {agentabi-0.2.0 → agentabi-0.2.2/src/agentabi.egg-info}/PKG-INFO +10 -5
  17. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/SOURCES.txt +5 -0
  18. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/requires.txt +2 -2
  19. agentabi-0.2.2/tests/test_middleware.py +512 -0
  20. {agentabi-0.2.0 → agentabi-0.2.2}/LICENSE +0 -0
  21. {agentabi-0.2.0 → agentabi-0.2.2}/setup.cfg +0 -0
  22. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/auto_detect.py +0 -0
  23. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/__init__.py +0 -0
  24. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/claude_sdk.py +0 -0
  25. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/gemini_sdk.py +0 -0
  26. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/registry.py +0 -0
  27. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/py.typed +0 -0
  28. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/__init__.py +0 -0
  29. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/__init__.py +0 -0
  30. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/capabilities.py +0 -0
  31. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/events.py +0 -0
  32. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/helpers.py +0 -0
  33. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/permissions.py +0 -0
  34. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/session.py +0 -0
  35. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/task.py +0 -0
  36. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/type_guards.py +0 -0
  37. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/dependency_links.txt +0 -0
  38. {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/top_level.txt +0 -0
  39. {agentabi-0.2.0 → agentabi-0.2.2}/tests/test_ir_types.py +0 -0
  40. {agentabi-0.2.0 → agentabi-0.2.2}/tests/test_session.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentabi
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: Unified async Python API for driving coding agent CLIs — Claude Code, Codex, Gemini CLI, and more
5
5
  Author-email: Peng Ding <pding.dp@foxmail.com>
6
6
  License-Expression: MIT
@@ -38,14 +38,19 @@ Provides-Extra: dev
38
38
  Requires-Dist: pytest>=7.0.0; extra == "dev"
39
39
  Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
40
40
  Requires-Dist: pytest-cov; extra == "dev"
41
- Requires-Dist: ruff>=0.15.0; extra == "dev"
42
- Requires-Dist: ty>=0.0.31; extra == "dev"
41
+ Requires-Dist: ruff==0.15.20; extra == "dev"
42
+ Requires-Dist: ty==0.0.54; extra == "dev"
43
43
  Requires-Dist: build>=1.2.2.post1; extra == "dev"
44
44
  Requires-Dist: twine>=6.1.0; extra == "dev"
45
45
  Dynamic: license-file
46
46
 
47
47
  # agentabi
48
48
 
49
+ [![PyPI version](https://img.shields.io/pypi/v/agentabi?color=green)](https://pypi.org/project/agentabi/)
50
+ [![GitHub release](https://img.shields.io/github/v/release/Oaklight/agentabi?color=green)](https://github.com/Oaklight/agentabi/releases/latest)
51
+ [![CI](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml/badge.svg)](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml)
52
+ [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
53
+
49
54
  Unified interface layer for agentic coding CLIs.
50
55
 
51
56
  One interface. Any coding agent.
@@ -138,10 +143,10 @@ print(caps["supports_streaming"]) # True
138
143
 
139
144
  ```
140
145
  agentabi → Agent CLI unified interface → like an OS ABI
141
- llmir → LLM API format conversion → like a compiler IR
146
+ llm-rosetta → LLM API format conversion → like a compiler IR
142
147
  ```
143
148
 
144
- - [llmir](https://github.com/Oaklight/llmir) — LLM Intermediate Representation for converting between LLM provider API formats (OpenAI, Anthropic, Google)
149
+ - [llm-rosetta](https://github.com/Oaklight/llm-rosetta) — LLM API format conversion layer for translating between LLM provider API formats (OpenAI, Anthropic, Google)
145
150
 
146
151
  ## License
147
152
 
@@ -1,5 +1,10 @@
1
1
  # agentabi
2
2
 
3
+ [![PyPI version](https://img.shields.io/pypi/v/agentabi?color=green)](https://pypi.org/project/agentabi/)
4
+ [![GitHub release](https://img.shields.io/github/v/release/Oaklight/agentabi?color=green)](https://github.com/Oaklight/agentabi/releases/latest)
5
+ [![CI](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml/badge.svg)](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml)
6
+ [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
7
+
3
8
  Unified interface layer for agentic coding CLIs.
4
9
 
5
10
  One interface. Any coding agent.
@@ -92,10 +97,10 @@ print(caps["supports_streaming"]) # True
92
97
 
93
98
  ```
94
99
  agentabi → Agent CLI unified interface → like an OS ABI
95
- llmir → LLM API format conversion → like a compiler IR
100
+ llm-rosetta → LLM API format conversion → like a compiler IR
96
101
  ```
97
102
 
98
- - [llmir](https://github.com/Oaklight/llmir) — LLM Intermediate Representation for converting between LLM provider API formats (OpenAI, Anthropic, Google)
103
+ - [llm-rosetta](https://github.com/Oaklight/llm-rosetta) — LLM API format conversion layer for translating between LLM provider API formats (OpenAI, Anthropic, Google)
99
104
 
100
105
  ## License
101
106
 
@@ -58,8 +58,8 @@ dev = [
58
58
  "pytest>=7.0.0",
59
59
  "pytest-asyncio>=0.21.0",
60
60
  "pytest-cov",
61
- "ruff>=0.15.0",
62
- "ty>=0.0.31",
61
+ "ruff==0.15.20",
62
+ "ty==0.0.54",
63
63
  "build>=1.2.2.post1",
64
64
  "twine>=6.1.0",
65
65
  ]
@@ -19,6 +19,14 @@ Sync convenience:
19
19
  from agentabi import run_sync
20
20
  result = run_sync(prompt="List Python files", agent="claude_code")
21
21
 
22
+ Middleware:
23
+ from agentabi.middleware import LoggingMiddleware, TimeoutMiddleware
24
+
25
+ session = Session(
26
+ agent="claude_code",
27
+ middleware=[LoggingMiddleware(), TimeoutMiddleware(120)],
28
+ )
29
+
22
30
  Discovery:
23
31
  from agentabi import detect_agents, get_agent_capabilities
24
32
  agents = detect_agents()
@@ -26,6 +34,14 @@ Discovery:
26
34
  """
27
35
 
28
36
  from .auto_detect import detect_agents, get_agent_capabilities, get_default_agent
37
+ from .middleware import (
38
+ LoggingMiddleware,
39
+ Middleware,
40
+ StreamHandler,
41
+ TimeoutMiddleware,
42
+ UsageMeterMiddleware,
43
+ chain_middleware,
44
+ )
29
45
  from .providers.base import Provider
30
46
  from .providers.registry import AgentNotAvailable, get_provider
31
47
  from .session import Session, run_sync
@@ -53,7 +69,7 @@ from .types.ir import (
53
69
  UsageEvent,
54
70
  )
55
71
 
56
- __version__ = "0.2.0"
72
+ __version__ = "0.2.2"
57
73
 
58
74
  __all__ = [
59
75
  # Consumer API
@@ -92,4 +108,11 @@ __all__ = [
92
108
  "PermissionConfig",
93
109
  "PermissionLevel",
94
110
  "PermissionRequest",
111
+ # Middleware
112
+ "Middleware",
113
+ "StreamHandler",
114
+ "chain_middleware",
115
+ "LoggingMiddleware",
116
+ "TimeoutMiddleware",
117
+ "UsageMeterMiddleware",
95
118
  ]
@@ -0,0 +1,119 @@
1
+ """
2
+ agentabi - Middleware Pipeline
3
+
4
+ Composable middleware for intercepting Session.stream() calls.
5
+ Middleware can inspect/modify TaskConfig, filter/transform IREvents,
6
+ and run side effects (logging, metering, timeouts).
7
+
8
+ Architecture:
9
+ Middleware follows the onion model (like ASGI). Each middleware
10
+ wraps a ``StreamHandler`` and returns a new ``StreamHandler``.
11
+ The first middleware in the list is the outermost wrapper.
12
+
13
+ Quick start::
14
+
15
+ from agentabi.middleware import LoggingMiddleware, TimeoutMiddleware
16
+
17
+ session = Session(
18
+ agent="claude_code",
19
+ middleware=[LoggingMiddleware(), TimeoutMiddleware(120)],
20
+ )
21
+
22
+ Custom middleware::
23
+
24
+ from agentabi.middleware import Middleware, StreamHandler
25
+
26
+ class MyMiddleware:
27
+ def __call__(self, handler: StreamHandler) -> StreamHandler:
28
+ async def wrapper(task):
29
+ task = {**task, "prompt": task["prompt"].strip()}
30
+ async for event in handler(task):
31
+ yield event
32
+ return wrapper
33
+ """
34
+
35
+ from __future__ import annotations
36
+
37
+ from collections.abc import AsyncIterator, Sequence
38
+ from typing import Callable
39
+
40
+ from ..types.ir.events import IREvent
41
+ from ..types.ir.task import TaskConfig
42
+ from .logging import LoggingMiddleware
43
+ from .timeout import TimeoutMiddleware
44
+ from .usage import UsageMeterMiddleware
45
+
46
+ # ============================================================================
47
+ # Core type aliases
48
+ # ============================================================================
49
+
50
+ StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
51
+ """A callable that takes a TaskConfig and returns an async iterator of IREvents.
52
+
53
+ This is the fundamental unit of the middleware pipeline. The provider's
54
+ ``stream()`` method is the innermost StreamHandler; each middleware wraps
55
+ it with additional behavior.
56
+ """
57
+
58
+ Middleware = Callable[[StreamHandler], StreamHandler]
59
+ """A callable that wraps a StreamHandler with additional behavior.
60
+
61
+ Middleware receives the next handler in the chain and returns a new
62
+ handler. The wrapper can:
63
+
64
+ - Modify TaskConfig before passing to the next handler (copy-on-mutate
65
+ recommended: ``task = {**task, "key": new_value}``)
66
+ - Filter or transform yielded IREvents
67
+ - Run side effects (logging, timing, metering)
68
+ - Short-circuit the pipeline (e.g., timeout)
69
+ """
70
+
71
+
72
+ # ============================================================================
73
+ # Pipeline chaining
74
+ # ============================================================================
75
+
76
+
77
+ def chain_middleware(
78
+ handler: StreamHandler,
79
+ middleware: Sequence[Middleware],
80
+ ) -> StreamHandler:
81
+ """Chain middleware around a base handler using the onion model.
82
+
83
+ The first middleware in the sequence becomes the outermost wrapper.
84
+ When the pipeline executes, control flows inward through each
85
+ middleware, reaches the base handler, then flows back outward.
86
+
87
+ Args:
88
+ handler: The base StreamHandler (typically provider.stream).
89
+ middleware: Sequence of middleware to apply. Empty sequence
90
+ returns the handler unchanged.
91
+
92
+ Returns:
93
+ A new StreamHandler with all middleware applied.
94
+
95
+ Example::
96
+
97
+ pipeline = chain_middleware(
98
+ provider.stream,
99
+ [LoggingMiddleware(), TimeoutMiddleware(30)],
100
+ )
101
+ async for event in pipeline(task):
102
+ ...
103
+ """
104
+ for mw in reversed(middleware):
105
+ handler = mw(handler)
106
+ return handler
107
+
108
+
109
+ __all__ = [
110
+ # Core types
111
+ "StreamHandler",
112
+ "Middleware",
113
+ # Pipeline
114
+ "chain_middleware",
115
+ # Built-in middleware
116
+ "LoggingMiddleware",
117
+ "TimeoutMiddleware",
118
+ "UsageMeterMiddleware",
119
+ ]
@@ -0,0 +1,112 @@
1
+ """
2
+ agentabi - Logging Middleware
3
+
4
+ Logs task execution and IR events at configurable levels.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import logging
10
+ import time
11
+ from collections.abc import AsyncIterator
12
+ from typing import Callable
13
+
14
+ from ..types.ir.events import IREvent
15
+ from ..types.ir.task import TaskConfig
16
+
17
+ StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
18
+
19
+ _DEFAULT_LOGGER = logging.getLogger("agentabi.middleware")
20
+
21
+
22
+ class LoggingMiddleware:
23
+ """Middleware that logs task execution and IR events.
24
+
25
+ Logs task start/end at INFO level and individual events at DEBUG
26
+ level. Uses the ``agentabi.middleware`` logger by default.
27
+
28
+ Args:
29
+ logger: Logger instance to use. Defaults to ``agentabi.middleware``.
30
+ level: Log level for event-level messages. Defaults to DEBUG.
31
+ log_events: Whether to log individual IR events. Defaults to True.
32
+ log_event_types: If set, only log events whose ``type`` is in
33
+ this set. Defaults to None (log all events).
34
+
35
+ Example::
36
+
37
+ import logging
38
+ logging.basicConfig(level=logging.DEBUG)
39
+
40
+ session = Session(
41
+ agent="claude_code",
42
+ middleware=[LoggingMiddleware()],
43
+ )
44
+ """
45
+
46
+ def __init__(
47
+ self,
48
+ *,
49
+ logger: logging.Logger | None = None,
50
+ level: int = logging.DEBUG,
51
+ log_events: bool = True,
52
+ log_event_types: set[str] | None = None,
53
+ ) -> None:
54
+ self._logger = logger or _DEFAULT_LOGGER
55
+ self._level = level
56
+ self._log_events = log_events
57
+ self._log_event_types = log_event_types
58
+
59
+ def __call__(self, handler: StreamHandler) -> StreamHandler:
60
+ """Wrap a StreamHandler with logging."""
61
+ logger = self._logger
62
+ level = self._level
63
+ log_events = self._log_events
64
+ log_event_types = self._log_event_types
65
+
66
+ async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
67
+ prompt_preview = task["prompt"][:80]
68
+ agent = task.get("agent", "unknown")
69
+ logger.info("Task started: agent=%s prompt=%r", agent, prompt_preview)
70
+
71
+ start = time.monotonic()
72
+ event_count = 0
73
+ failed = False
74
+
75
+ try:
76
+ async for event in handler(task):
77
+ event_count += 1
78
+ if log_events:
79
+ etype = event.get("type", "unknown")
80
+ if log_event_types is None or etype in log_event_types:
81
+ logger.log(level, "Event #%d: type=%s", event_count, etype)
82
+ yield event
83
+ except Exception:
84
+ failed = True
85
+ raise
86
+ finally:
87
+ elapsed = time.monotonic() - start
88
+ if failed:
89
+ logger.error(
90
+ "Task failed: agent=%s elapsed=%.2fs events=%d",
91
+ agent,
92
+ elapsed,
93
+ event_count,
94
+ )
95
+ else:
96
+ logger.info(
97
+ "Task completed: agent=%s elapsed=%.2fs events=%d",
98
+ agent,
99
+ elapsed,
100
+ event_count,
101
+ )
102
+
103
+ return wrapper
104
+
105
+ def __repr__(self) -> str:
106
+ return (
107
+ f"LoggingMiddleware(logger={self._logger.name!r}, "
108
+ f"level={self._level}, log_events={self._log_events})"
109
+ )
110
+
111
+
112
+ __all__ = ["LoggingMiddleware"]
@@ -0,0 +1,84 @@
1
+ """
2
+ agentabi - Timeout Middleware
3
+
4
+ Enforces a wall-clock timeout on streaming task execution.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import asyncio
10
+ import time
11
+ from collections.abc import AsyncIterator
12
+ from typing import Callable
13
+
14
+ from ..types.ir.events import IREvent
15
+ from ..types.ir.task import TaskConfig
16
+
17
+ StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
18
+
19
+
20
+ class TimeoutMiddleware:
21
+ """Middleware that enforces a wall-clock timeout on task execution.
22
+
23
+ Wraps each ``__anext__`` call on the event stream with
24
+ ``asyncio.wait_for``, so the timeout fires even if the stream
25
+ stalls mid-iteration (not just between events).
26
+
27
+ Raises ``asyncio.TimeoutError`` if the total elapsed time exceeds
28
+ the configured timeout.
29
+
30
+ Args:
31
+ timeout: Maximum wall-clock seconds for the entire stream.
32
+ Must be positive.
33
+
34
+ Example::
35
+
36
+ session = Session(
37
+ agent="claude_code",
38
+ middleware=[TimeoutMiddleware(120)],
39
+ )
40
+
41
+ try:
42
+ result = await session.run(prompt="Long task")
43
+ except asyncio.TimeoutError:
44
+ print("Task timed out after 120s")
45
+
46
+ Note:
47
+ The timeout is per-stream invocation, not cumulative across
48
+ multiple calls. Each ``stream()`` or ``run()`` gets a fresh
49
+ timeout budget.
50
+ """
51
+
52
+ def __init__(self, timeout: float) -> None:
53
+ if timeout <= 0:
54
+ raise ValueError(f"timeout must be positive, got {timeout}")
55
+ self._timeout = timeout
56
+
57
+ def __call__(self, handler: StreamHandler) -> StreamHandler:
58
+ """Wrap a StreamHandler with timeout enforcement."""
59
+ timeout = self._timeout
60
+
61
+ async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
62
+ start = time.monotonic()
63
+ aiter = handler(task).__aiter__()
64
+
65
+ while True:
66
+ elapsed = time.monotonic() - start
67
+ remaining = timeout - elapsed
68
+ if remaining <= 0:
69
+ raise asyncio.TimeoutError(
70
+ f"Task exceeded timeout of {timeout}s (elapsed: {elapsed:.1f}s)"
71
+ )
72
+ try:
73
+ event = await asyncio.wait_for(aiter.__anext__(), timeout=remaining)
74
+ except StopAsyncIteration:
75
+ break
76
+ yield event
77
+
78
+ return wrapper
79
+
80
+ def __repr__(self) -> str:
81
+ return f"TimeoutMiddleware(timeout={self._timeout})"
82
+
83
+
84
+ __all__ = ["TimeoutMiddleware"]
@@ -0,0 +1,134 @@
1
+ """
2
+ agentabi - Usage Meter Middleware
3
+
4
+ Tracks cumulative token usage and cost across multiple task executions.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import threading
10
+ from collections.abc import AsyncIterator
11
+ from typing import Callable
12
+
13
+ from typing_extensions import TypedDict
14
+
15
+ from ..types.ir.events import IREvent
16
+ from ..types.ir.task import TaskConfig
17
+
18
+ StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
19
+
20
+
21
+ class UsageSummary(TypedDict, total=False):
22
+ """Cumulative usage statistics across multiple task executions."""
23
+
24
+ call_count: int
25
+ total_input_tokens: int
26
+ total_output_tokens: int
27
+ total_cache_read_tokens: int
28
+ total_cache_creation_tokens: int
29
+ total_tokens: int
30
+ total_cost_usd: float
31
+
32
+
33
+ class UsageMeterMiddleware:
34
+ """Middleware that accumulates token usage and cost across calls.
35
+
36
+ Watches for ``usage`` events in the stream and accumulates token
37
+ counts and cost. The accumulated totals are available via the
38
+ ``summary`` property at any time.
39
+
40
+ Thread-safe: uses a lock for counter updates so the summary can
41
+ be read from any thread.
42
+
43
+ Example::
44
+
45
+ meter = UsageMeterMiddleware()
46
+ session = Session(agent="claude_code", middleware=[meter])
47
+
48
+ await session.run(prompt="Task 1")
49
+ await session.run(prompt="Task 2")
50
+
51
+ print(meter.summary)
52
+ # {'call_count': 2, 'total_input_tokens': 1500, ...}
53
+
54
+ meter.reset()
55
+ """
56
+
57
+ def __init__(self) -> None:
58
+ self._lock = threading.Lock()
59
+ self._call_count = 0
60
+ self._input_tokens = 0
61
+ self._output_tokens = 0
62
+ self._cache_read_tokens = 0
63
+ self._cache_creation_tokens = 0
64
+ self._total_tokens = 0
65
+ self._cost_usd = 0.0
66
+
67
+ @property
68
+ def summary(self) -> UsageSummary:
69
+ """Return cumulative usage statistics.
70
+
71
+ Returns:
72
+ A UsageSummary dict with all accumulated counters.
73
+ """
74
+ with self._lock:
75
+ result: UsageSummary = {
76
+ "call_count": self._call_count,
77
+ "total_input_tokens": self._input_tokens,
78
+ "total_output_tokens": self._output_tokens,
79
+ "total_tokens": self._total_tokens,
80
+ "total_cost_usd": self._cost_usd,
81
+ }
82
+ if self._cache_read_tokens:
83
+ result["total_cache_read_tokens"] = self._cache_read_tokens
84
+ if self._cache_creation_tokens:
85
+ result["total_cache_creation_tokens"] = self._cache_creation_tokens
86
+ return result
87
+
88
+ def reset(self) -> None:
89
+ """Reset all accumulated counters to zero."""
90
+ with self._lock:
91
+ self._call_count = 0
92
+ self._input_tokens = 0
93
+ self._output_tokens = 0
94
+ self._cache_read_tokens = 0
95
+ self._cache_creation_tokens = 0
96
+ self._total_tokens = 0
97
+ self._cost_usd = 0.0
98
+
99
+ def __call__(self, handler: StreamHandler) -> StreamHandler:
100
+ """Wrap a StreamHandler with usage metering."""
101
+
102
+ async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
103
+ with self._lock:
104
+ self._call_count += 1
105
+
106
+ async for event in handler(task):
107
+ if event.get("type") == "usage":
108
+ usage = event.get("usage", {})
109
+ cost = event.get("cost_usd")
110
+ with self._lock:
111
+ self._input_tokens += usage.get("input_tokens", 0)
112
+ self._output_tokens += usage.get("output_tokens", 0)
113
+ self._cache_read_tokens += usage.get("cache_read_tokens", 0)
114
+ self._cache_creation_tokens += usage.get(
115
+ "cache_creation_tokens", 0
116
+ )
117
+ self._total_tokens += usage.get("total_tokens", 0)
118
+ if cost is not None:
119
+ self._cost_usd += cost
120
+ yield event
121
+
122
+ return wrapper
123
+
124
+ def __repr__(self) -> str:
125
+ return (
126
+ f"UsageMeterMiddleware(calls={self._call_count}, "
127
+ f"tokens={self._total_tokens}, cost=${self._cost_usd:.4f})"
128
+ )
129
+
130
+
131
+ __all__ = [
132
+ "UsageMeterMiddleware",
133
+ "UsageSummary",
134
+ ]
@@ -7,13 +7,14 @@ abstraction layer between Session and agent CLIs/SDKs.
7
7
 
8
8
  from __future__ import annotations
9
9
 
10
+ import asyncio
10
11
  from collections.abc import AsyncIterator
11
12
  from typing import Any
12
13
 
13
14
  from typing_extensions import Protocol, runtime_checkable
14
15
 
15
16
  from ..types.ir.capabilities import AgentCapabilities
16
- from ..types.ir.events import IREvent, UsageInfo
17
+ from ..types.ir.events import ErrorEvent, IREvent, UsageInfo
17
18
  from ..types.ir.session import SessionResult, SessionStatus
18
19
  from ..types.ir.task import TaskConfig
19
20
 
@@ -71,6 +72,54 @@ class Provider(Protocol):
71
72
  ...
72
73
 
73
74
 
75
+ async def collect_subprocess_errors(
76
+ proc: asyncio.subprocess.Process,
77
+ *,
78
+ timed_out: bool = False,
79
+ timeout_seconds: float | None = None,
80
+ ) -> list[ErrorEvent]:
81
+ """Collect stderr and exit-code errors after a subprocess finishes.
82
+
83
+ Call after the stdout reading loop ends and ``await proc.wait()``
84
+ has been called so that ``proc.returncode`` is set.
85
+
86
+ Args:
87
+ proc: The completed subprocess.
88
+ timed_out: Whether the process was killed due to timeout.
89
+ timeout_seconds: The configured timeout value (for the error message).
90
+
91
+ Returns:
92
+ A (possibly empty) list of ErrorEvents to yield.
93
+ """
94
+ events: list[ErrorEvent] = []
95
+ stderr_text = ""
96
+
97
+ if proc.stderr:
98
+ stderr_data = await proc.stderr.read()
99
+ if stderr_data:
100
+ stderr_text = stderr_data.decode().strip()
101
+ if stderr_text:
102
+ events.append(ErrorEvent(type="error", error=stderr_text))
103
+
104
+ if timed_out:
105
+ events.append(
106
+ ErrorEvent(
107
+ type="error",
108
+ error=f"Process timed out after {timeout_seconds}s",
109
+ is_fatal=True,
110
+ )
111
+ )
112
+ elif proc.returncode:
113
+ events.append(
114
+ ErrorEvent(
115
+ type="error",
116
+ error=f"Process exited with code {proc.returncode}",
117
+ )
118
+ )
119
+
120
+ return events
121
+
122
+
74
123
  class _RunState:
75
124
  """Mutable accumulator for default_run event aggregation."""
76
125
 
@@ -163,5 +212,6 @@ async def default_run(provider: Any, task: TaskConfig) -> SessionResult:
163
212
 
164
213
  __all__ = [
165
214
  "Provider",
215
+ "collect_subprocess_errors",
166
216
  "default_run",
167
217
  ]