axio 0.11.3__tar.gz → 0.12.0__tar.gz

{axio-0.11.3 → axio-0.12.0}/.gitignore

@@ -12,3 +12,4 @@ _build/
 .DS_Store
 *.sqlite
 *.db
+.axio/mcp.local.json

{axio-0.11.3 → axio-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: axio
-Version: 0.11.3
+Version: 0.12.0
 Summary: Minimal, streaming-first, protocol-driven foundation for LLM-powered agents
 Project-URL: Documentation, https://docs.axio-agent.com
 Project-URL: Homepage, https://github.com/mosquito/axio-agent

{axio-0.11.3 → axio-0.12.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "axio"
-version = "0.11.3"  # bumped to the git tag by CI on publish
+version = "0.12.0"  # bumped to the git tag by CI on publish
 description = "Minimal, streaming-first, protocol-driven foundation for LLM-powered agents"
 readme = "README.md"
 requires-python = ">=3.12"

{axio-0.11.3 → axio-0.12.0}/src/axio/agent.py

@@ -7,7 +7,7 @@ import dataclasses
 import json
 import logging
 import time
-from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Iterator
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping
 from contextlib import aclosing, suppress
 from dataclasses import dataclass, field, replace
 from datetime import datetime
@@ -45,11 +45,15 @@ from .events import (
     VideoOutput,
 )
 from .exceptions import StreamError
+from .instructions import LOAD_INSTRUCTION_TOOL, Instruction, InstructionState, LoadedInstruction
 from .messages import Message
 from .models import Capability
+from .permission import PermissionGuard
+from .plugins import ToolProvider, collect_tools
 from .selector import ToolSelector
 from .stream import AgentStream
 from .tool import Tool
+from .tool_access import Access
 from .transport import CompletionTransport
 from .types import StopReason, Usage
 
@@ -281,9 +285,100 @@ class Agent:
     selector: ToolSelector | None = field(default=None)
     max_iterations: int = field(default=50)
     last_iteration_message: Message | None = field(default=None)
+    before_turn: Callable[[int], Awaitable[None]] | None = field(default=None, kw_only=True)
+    before_tools: Callable[[list[ToolUseBlock], int], Awaitable[list[ToolResultBlock]]] | None = field(
+        default=None, kw_only=True
+    )
+    providers: list[ToolProvider] = field(default_factory=list, kw_only=True)
+    instructions: list[Instruction] = field(default_factory=list, kw_only=True)
+    required_instructions: tuple[str, ...] = field(default=(), kw_only=True)
+    before_turn_hooks: tuple[Callable[[int], Awaitable[None]], ...] = field(default=(), kw_only=True)
+    before_tools_hooks: tuple[Callable[[list[ToolUseBlock], int], Awaitable[list[ToolResultBlock]]], ...] = field(
+        default=(), kw_only=True
+    )
+    system_contributions: tuple[Callable[[], str], ...] = field(default=(), kw_only=True)
+    tool_guards: tuple[PermissionGuard, ...] = field(default=(), kw_only=True)
+    _instruction_state: InstructionState = field(default_factory=InstructionState, init=False, repr=False)
+    _instruction_tool: Tool[Any] | None = field(default=None, init=False, repr=False)
+    _iteration_tools: list[Tool[Any]] | None = field(default=None, init=False, repr=False)
+    _guarded_tools: dict[int, tuple[Tool[Any], tuple[PermissionGuard, ...], Tool[Any]]] = field(
+        default_factory=dict, init=False, repr=False
+    )
+
+    def _catalog(self) -> list[Instruction]:
+        return [*self.instructions, *(item for provider in self.providers for item in provider.get_instructions())]
+
+    def _base_tools(self) -> list[Tool[Any]]:
+        tools = collect_tools(self.providers, tools=self.tools)
+        if not self._catalog():
+            return tools
+        if self._instruction_tool is None:
+            self._instruction_tool = Tool(
+                name=LOAD_INSTRUCTION_TOOL, handler=self._load_instruction_tool, access=self._instruction_access
+            )
+        return [*tools, self._instruction_tool]
+
+    def get_tools(self) -> list[Tool[Any]]:
+        """Return collected tools before selection, including activated instructions."""
+        self._instruction_state.synchronize(self._catalog())
+        return self._apply_tool_guards(
+            self._instruction_state.collect_tools(self._base_tools(), self._instruction_state.loaded.values())
+        )
+
+    def _instruction_access(self, values: Mapping[str, Any]) -> tuple[Access, ...]:
+        id = values.get("id")
+        for instruction in self._catalog():
+            if instruction.id == id:
+                return instruction.load_access()
+        raise ValueError(f"Unknown instruction: {id!r}")
+
+    async def load_instruction(self, id: str) -> LoadedInstruction:
+        """Load provider content and stage it for the next model iteration."""
+        self._instruction_state.synchronize(self._catalog())
+        return await self._instruction_state.load(id, self._base_tools())
+
+    async def _load_instruction_tool(self, id: str) -> str:
+        """Load a listed instruction and make its associated tools available on the next turn."""
+        return (await self.load_instruction(id)).text
+
+    async def _prepare_instructions(self) -> list[Tool[Any]]:
+        self._instruction_state.synchronize(self._catalog())
+        try:
+            for id in self.required_instructions:
+                await self.load_instruction(id)
+            return self._instruction_state.commit(self._base_tools())
+        except BaseException:
+            self._instruction_state.pending.clear()
+            raise
+
+    def _effective_system(self) -> str:
+        parts = (self.system, *(part() for part in self.system_contributions), self._instruction_state.render())
+        return "\n\n".join(filter(None, parts))
+
+    def _apply_tool_guards(self, tools: list[Tool[Any]]) -> list[Tool[Any]]:
+        """Append application guards to every tool, including newly loaded tools."""
+        if not self.tool_guards:
+            self._guarded_tools.clear()
+            return tools
+        cached: dict[int, tuple[Tool[Any], tuple[PermissionGuard, ...], Tool[Any]]] = {}
+        result: list[Tool[Any]] = []
+        for tool in tools:
+            entry = self._guarded_tools.get(id(tool))
+            if entry is None or entry[0] is not tool or entry[1] != self.tool_guards:
+                guards = (*tool.guards, *(guard for guard in self.tool_guards if guard not in tool.guards))
+                guarded = replace(tool, guards=guards)
+                # Guard wrappers share the provider tool's concurrency budget across agents.
+                object.__setattr__(guarded, "_semaphore", tool._semaphore)
+                entry = (tool, self.tool_guards, guarded)
+            cached[id(tool)] = entry
+            result.append(entry[2])
+        self._guarded_tools = cached
+        return result
 
     def copy(self, **overrides: Any) -> Self:
-        """Return a new Agent with *overrides* applied."""
+        """Return a new Agent with fresh instruction state and *overrides* applied."""
+        overrides.setdefault("instructions", list(self.instructions))
+        overrides.setdefault("providers", list(self.providers))
         return dataclasses.replace(self, **overrides)
 
     def run_stream(self, user_message: str, context: ContextStore) -> AgentStream:
@@ -315,11 +410,54 @@ class Agent:
 
     async def dispatch_tools(self, blocks: list[ToolUseBlock], iteration: int) -> list[ToolResultBlock]:
         logger.info("Dispatching %d tool(s): %r", len(blocks), blocks)
-        results = list(await asyncio.gather(*[self._call_one(b) for b in blocks]))
+        runnable, deferred = await self._preflight_tools(blocks, iteration)
+        results = [*deferred, *await asyncio.gather(*[self._call_one(b) for b in runnable])]
         error_count = sum(1 for r in results if r.is_error)
         logger.info("Tools complete: %d total, %d errors", len(results), error_count)
         return results
 
+    async def _preflight_tools(
+        self, blocks: list[ToolUseBlock], iteration: int
+    ) -> tuple[list[ToolUseBlock], list[ToolResultBlock]]:
+        """Get replacement results before any handler in this batch can run.
+
+        The callback can defer all calls or a subset. Each returned result must
+        refer to a distinct call in this batch. Exceptions fail the whole batch;
+        cancellation propagates to the normal interrupted-result handling.
+        """
+        hooks = (*((self.before_tools,) if self.before_tools is not None else ()), *self.before_tools_hooks)
+        if not hooks:
+            return blocks, []
+        try:
+            results: list[ToolResultBlock] = []
+            remaining = blocks
+            for hook in hooks:
+                if not remaining:
+                    break
+                replacement = await hook(remaining, iteration)
+                ids = [result.tool_use_id for result in replacement]
+                deferred_ids = set(ids)
+                if len(deferred_ids) != len(ids) or not deferred_ids <= {block.id for block in remaining}:
+                    raise ValueError("Tool preflight returned duplicate or unknown call IDs")
+                results.extend(replacement)
+                remaining = [block for block in remaining if block.id not in deferred_ids]
+            if self._instruction_state.pending:
+                return [], [
+                    ToolResultBlock(
+                        tool_use_id=block.id,
+                        content="Instructions changed. Review the updated instructions before retrying this batch.",
+                        is_error=True,
+                    )
+                    for block in blocks
+                ]
+        except Exception as error:
+            self._instruction_state.pending.clear()
+            return [], [
+                ToolResultBlock(tool_use_id=block.id, content=f"Tool preflight failed: {error}", is_error=True)
+                for block in blocks
+            ]
+        return remaining, results
+
     async def _dispatch_tools_streaming(
         self,
         blocks: list[ToolUseBlock],
@@ -328,6 +466,7 @@ class Agent:
     ) -> list[ToolResultBlock]:
         """Like dispatch_tools but pushes ToolOutputDelta events for streaming tools."""
         logger.info("Dispatching %d tool(s) with streaming: %r", len(blocks), blocks)
+        runnable, deferred = await self._preflight_tools(blocks, iteration)
 
         async def _run_one(block: ToolUseBlock) -> ToolResultBlock:
             tool = self._find_tool(block.name)
@@ -345,13 +484,24 @@ class Agent:
                 return ToolResultBlock(tool_use_id=block.id, content=str(exc), is_error=True)
             return ToolResultBlock(tool_use_id=block.id, content=tool.format_stream_result(chunks))
 
-        results = list(await asyncio.gather(*[_run_one(b) for b in blocks]))
+        results = [*deferred, *await asyncio.gather(*[_run_one(b) for b in runnable])]
         error_count = sum(1 for r in results if r.is_error)
         logger.info("Tools complete: %d total, %d errors", len(results), error_count)
         return results
 
     def _find_tool(self, name: str) -> Tool[Any] | None:
-        for tool in self.tools:
+        return self.get_tool(name)
+
+    def get_tool(self, name: str) -> Tool[Any] | None:
+        """Return the dispatchable tool, including its access declaration and guards.
+
+        During a run this uses the current iteration's selected tool set.
+        Preflight hooks can inspect effects without accessing a second registry.
+        """
+        tools = self._iteration_tools
+        if tools is None:
+            tools = self.get_tools()
+        for tool in tools:
             if tool.name == name:
                 return tool
         return None
@@ -679,6 +829,18 @@ class Agent:
 
         try:
             for iteration in range(1, self.max_iterations + 1):
+                try:
+                    if self.before_turn is not None:
+                        await self.before_turn(iteration)
+                    for hook in self.before_turn_hooks:
+                        await hook(iteration)
+                    tools = self._apply_tool_guards(await self._prepare_instructions())
+                    system = self._effective_system()
+                except Exception as exc:
+                    yield Error(exception=exc)
+                    yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
+                    session_end_emitted = True
+                    return
                 history = await context.get_history()
                 logger.info("Iteration %d, history length=%d", iteration, len(history))
                 effective_history = (
@@ -691,10 +853,11 @@ class Agent:
                 if model_caps is not None and Capability.tool_use not in model_caps:
                     active_tools: list[Tool[Any]] = []
                 else:
-                    active_tools = list(await self._select_tools(effective_history, self.tools))
+                    active_tools = list(await self._select_tools(effective_history, tools))
+                self._iteration_tools = active_tools
 
                 turn = _Turn()
-                stream = self.transport.stream(effective_history, active_tools, self.system)
+                stream = self.transport.stream(effective_history, active_tools, system)
                 try:
                     # `_consume` closes the stream in its `finally`, which runs only on close.
                     # Left to the collector, a caller that walked away held the response open.
@@ -828,3 +991,6 @@ class Agent:
             if not session_end_emitted:
                 yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
             raise
+        finally:
+            self._iteration_tools = None
+            self._instruction_state.pending.clear()

axio-0.12.0/src/axio/credentials/__init__.py

@@ -0,0 +1,15 @@
+"""Typed credential storage shared by transport providers."""
+
+from __future__ import annotations
+
+from .codec import CredentialCodec as CredentialCodec
+from .codec import record_id as record_id
+from .codec import record_key as record_key
+from .stores import Credentials as Credentials
+from .stores import FileCredentialStore as FileCredentialStore
+from .stores import SQLiteCredentialStore as SQLiteCredentialStore
+from .stores import default_credentials_path as default_credentials_path
+from .types import CredentialConflict as CredentialConflict
+from .types import CredentialError as CredentialError
+from .types import CredentialStore as CredentialStore
+from .types import complete_operation as complete_operation

axio-0.12.0/src/axio/credentials/codec.py

@@ -0,0 +1,148 @@
+"""Validate and serialize typed credential dataclasses without provider codecs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import types
+from dataclasses import asdict, fields, is_dataclass
+from enum import Enum
+from typing import Annotated, Any, Literal, Union, cast, get_args, get_origin, get_type_hints
+
+from .types import CredentialError
+
+
+def record_id(record: object) -> str:
+    """Find the one dataclass field marked with metadata={\"credential_key\": True}."""
+    keys: list[object] = []
+
+    def visit(value: object) -> None:
+        if not is_dataclass(value) or isinstance(value, type):
+            return
+        for item in fields(value):
+            child = getattr(value, item.name)
+            if item.metadata.get("credential_key"):
+                keys.append(child)
+            else:
+                visit(child)
+
+    visit(record)
+    if len(keys) != 1 or not isinstance(keys[0], str) or not keys[0]:
+        raise CredentialError("Credentials require one nonempty string key.")
+    return keys[0]
+
+
+def record_key(key: str) -> str:
+    """Hash record identifiers to keep them out of filesystem paths."""
+    if not isinstance(key, str) or not key:
+        raise CredentialError("A nonempty credential key is required.")
+    # Keep the original prefix so existing stores use the same process locks.
+    return "account-" + hashlib.sha256(key.encode()).hexdigest()
+
+
+def json_value(value: Any) -> Any:
+    """Convert dataclass containers to JSON values without stringifying secrets."""
+    if isinstance(value, Enum):
+        return json_value(value.value)
+    if isinstance(value, dict):
+        if any(not isinstance(key, str) for key in value):
+            raise ValueError
+        return {key: json_value(item) for key, item in value.items()}
+    if isinstance(value, (set, frozenset)):
+        return sorted((json_value(item) for item in value), key=lambda item: json.dumps(item, sort_keys=True))
+    if isinstance(value, (tuple, list)):
+        return [json_value(item) for item in value]
+    if value is None or type(value) in (str, int, bool):
+        return value
+    if type(value) is float and math.isfinite(value):
+        return value
+    raise ValueError
+
+
+def restore(value: Any, annotation: Any) -> Any:
+    """Restore declared JSON-compatible types without coercing invalid values."""
+    origin, args = get_origin(annotation), get_args(annotation)
+    if origin is Annotated:
+        return restore(value, args[0])
+    if origin in (types.UnionType, Union):
+        for member in args:
+            try:
+                return restore(value, member)
+            except (TypeError, ValueError):
+                pass
+        raise ValueError
+    if origin is Literal:
+        if any(type(value) is type(choice) and value == choice for choice in args):
+            return value
+        raise ValueError
+    if annotation is type(None) and value is None:
+        return None
+    if annotation in (str, int, bool) and type(value) is annotation:
+        return value
+    if annotation is float and type(value) in (int, float) and math.isfinite(value):
+        return float(value)
+    if isinstance(annotation, type) and issubclass(annotation, Enum):
+        for member in annotation:
+            if type(value) is type(member.value) and value == member.value:
+                return member
+        raise ValueError
+    if origin in (list, set, frozenset, tuple) and isinstance(value, list):
+        if origin is tuple and (not args or args[-1] is not Ellipsis):
+            if len(value) != len(args):
+                raise ValueError
+            return tuple(restore(item, kind) for item, kind in zip(value, args, strict=True))
+        return origin(restore(item, args[0]) for item in value)
+    if origin is dict and args[0] is str and isinstance(value, dict):
+        if any(not isinstance(key, str) for key in value):
+            raise ValueError
+        return {key: restore(item, args[1]) for key, item in value.items()}
+    if isinstance(annotation, type) and is_dataclass(annotation) and isinstance(value, dict):
+        declared = {item.name for item in fields(annotation) if item.init}
+        if value.keys() != declared:
+            raise ValueError
+        hints = get_type_hints(annotation, include_extras=True)
+        return annotation(**{name: restore(value[name], hints[name]) for name in declared})
+    raise ValueError
+
+
+class CredentialCodec[T]:
+    """Versioned dataclass records with safe errors and strict type checks."""
+
+    def __init__(self, record_type: type[T]) -> None:
+        if not isinstance(record_type, type) or not is_dataclass(record_type):
+            raise TypeError("Credential records must be dataclasses.")
+        self.record_type = record_type
+
+    def from_dict(self, data: Any, key: str | None = None) -> T:
+        try:
+            if not isinstance(data, dict) or type(data.get("version")) is not int or data["version"] != 1:
+                raise ValueError
+            record = cast(
+                T, restore({name: value for name, value in data.items() if name != "version"}, self.record_type)
+            )
+            identity = record_id(record)
+            if key is not None and identity != key:
+                raise ValueError
+            return record
+        except (ValueError, TypeError, KeyError, OverflowError, RecursionError):
+            raise CredentialError("Invalid stored credential record.") from None
+
+    def to_dict(self, record: T) -> dict[str, Any]:
+        try:
+            if not isinstance(record, self.record_type) or not is_dataclass(record) or isinstance(record, type):
+                raise ValueError
+            data: dict[str, Any] = {"version": 1, **json_value(asdict(record))}
+            self.from_dict(data, record_id(record))
+            return data
+        except (ValueError, TypeError, OverflowError, RecursionError):
+            raise CredentialError("Invalid credential record.") from None
+
+    def decode(self, raw: str, key: str | None = None) -> T:
+        try:
+            return self.from_dict(json.loads(raw), key)
+        except (ValueError, TypeError, RecursionError):
+            raise CredentialError("Invalid stored credential record.") from None
+
+    def encode(self, record: T) -> str:
+        return json.dumps(self.to_dict(record), allow_nan=False, separators=(",", ":"))

axio-0.12.0/src/axio/credentials/files.py

@@ -0,0 +1,148 @@
+"""Protected files and locks for local credential backends."""
+
+from __future__ import annotations
+
+import asyncio
+import errno
+import os
+import stat
+import tempfile
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+from pathlib import Path
+
+from .types import CredentialError, complete_operation
+
+if os.name == "posix":
+    import fcntl
+
+
+def check_file(info: os.stat_result, *, directory: bool = False) -> None:
+    expected = stat.S_ISDIR(info.st_mode) if directory else stat.S_ISREG(info.st_mode)
+    protected = not (info.st_mode & 0o022) if directory else stat.S_IMODE(info.st_mode) == 0o600
+    if not expected or info.st_uid != os.getuid() or not protected:
+        raise CredentialError(
+            "Credential paths must be owned by this user, files mode 0600, directories not writable by others."
+        )
+    if not directory and info.st_nlink != 1:
+        raise CredentialError("Credential files must not have hard links.")
+
+
+def open_file(path: Path, flags: int) -> int:
+    fd = os.open(path, flags | os.O_NOFOLLOW | os.O_CLOEXEC | os.O_NONBLOCK, 0o600)
+    try:
+        check_file(os.fstat(fd))
+    except BaseException:
+        os.close(fd)
+        raise
+    return fd
+
+
+def sync_directory(path: Path) -> None:
+    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def make_directory(path: Path) -> None:
+    try:
+        path.mkdir(mode=0o700)
+    except FileNotFoundError:
+        make_directory(path.parent)
+        make_directory(path)
+    except FileExistsError:
+        pass
+    else:
+        sync_directory(path.parent)
+
+
+class ProtectedDirectory:
+    def __init__(self, path: Path) -> None:
+        if os.name != "posix":
+            raise CredentialError("Local credential stores require POSIX permissions and file locks.")
+        self.path = path.absolute()
+
+    def prepare(self) -> None:
+        try:
+            make_directory(self.path)
+            check_file(self.path.lstat(), directory=True)
+        except OSError:
+            raise CredentialError("Could not open the protected credential directory.") from None
+
+    def read(self, name: str) -> str | None:
+        self.prepare()
+        try:
+            fd = open_file(self.path / name, os.O_RDONLY)
+        except FileNotFoundError:
+            return None
+        except OSError:
+            raise CredentialError("Could not open the protected credential file.") from None
+        try:
+            with os.fdopen(fd, "r", encoding="utf-8") as stream:
+                return stream.read()
+        except (OSError, UnicodeError):
+            raise CredentialError("Could not read the credential file.") from None
+
+    def write(self, name: str, raw: str) -> None:
+        self.prepare()
+        temporary: str | None = None
+        try:
+            destination = self.path / name
+            try:
+                check_file(destination.lstat())
+            except FileNotFoundError:
+                pass
+            fd, temporary = tempfile.mkstemp(prefix=".credential-", dir=self.path)
+            with os.fdopen(fd, "w", encoding="utf-8") as stream:
+                stream.write(raw)
+                stream.flush()
+                os.fsync(stream.fileno())
+            os.replace(temporary, destination)
+            sync_directory(self.path)
+        except OSError:
+            raise CredentialError("Could not commit the credential file.") from None
+        finally:
+            if temporary is not None:
+                try:
+                    Path(temporary).unlink(missing_ok=True)
+                except OSError:
+                    pass
+
+    @asynccontextmanager
+    async def lock(self, key: str) -> AsyncIterator[None]:
+        await complete_operation(asyncio.to_thread(self.prepare))
+        try:
+            fd = open_file(self.path / (key + ".lock"), os.O_RDWR | os.O_CREAT)
+        except OSError:
+            raise CredentialError("Could not open the credential lock.") from None
+        try:
+            while True:
+                try:
+                    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    break
+                except OSError as exc:
+                    if exc.errno not in (errno.EACCES, errno.EAGAIN):
+                        raise CredentialError("Could not acquire the credential lock.") from None
+                    await asyncio.sleep(0.025)
+            yield
+        finally:
+            os.close(fd)
+
+    @contextmanager
+    def blocking_lock(self, key: str) -> Iterator[None]:
+        """Hold a short file-update lock from a worker thread."""
+        self.prepare()
+        try:
+            fd = open_file(self.path / (key + ".lock"), os.O_RDWR | os.O_CREAT)
+        except OSError:
+            raise CredentialError("Could not open the credential update lock.") from None
+        try:
+            try:
+                fcntl.flock(fd, fcntl.LOCK_EX)
+            except OSError:
+                raise CredentialError("Could not acquire the credential update lock.") from None
+            yield
+        finally:
+            os.close(fd)