iac-code 0.13.1__tar.gz → 0.14.0__tar.gz

{iac_code-0.13.1 → iac_code-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: iac_code
-Version: 0.13.1
+Version: 0.14.0
 Summary: Your AI-powered Infrastructure as Code assistant
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
@@ -14,6 +14,7 @@ Provides-Extra: a2a
 Provides-Extra: a2a-grpc
 Provides-Extra: a2a-redis
 Provides-Extra: a2a-signing
+Provides-Extra: agui
 Provides-Extra: diagram
 Provides-Extra: http
 License-File: LICENSE

{iac_code-0.13.1 → iac_code-0.14.0}/pyproject.toml

@@ -74,6 +74,12 @@ a2a-grpc = [
 a2a-redis = [
     "redis>=5.0.0",
 ]
+agui = [
+    "ag-ui-protocol==0.1.20",
+    "a2a-sdk[http-server,signing]>=1.0.2,<2",
+    "starlette>=0.39.0",
+    "uvicorn[standard]>=0.30.0",
+]
 diagram = ["termaid>=0.1; python_version >= '3.11'"]
 
 [dependency-groups]
@@ -146,6 +152,9 @@ default = true
 
 [tool.pytest.ini_options]
 timeout = 30
+markers = [
+    "integration: process-level integration tests that use local services only",
+]
 
 [tool.coverage.run]
 source_pkgs = ["iac_code"]

iac_code-0.14.0/src/iac_code/__init__.py

@@ -0,0 +1,2 @@
+__version__ = "0.14.0"
+__release_date__ = "2026-08-28"

{iac_code-0.13.1 → iac_code-0.14.0}/src/iac_code/a2a/app.py

@@ -39,9 +39,12 @@ from iac_code.a2a.projection import (
     resolve_a2a_public_path_roots,
     resolve_a2a_public_path_roots_for_data,
 )
+from iac_code.a2a.types import validate_protocol_id
 from iac_code.i18n import _
 from iac_code.pipeline.config import get_run_mode
 from iac_code.services.configuration_readiness import configuration_readiness
+from iac_code.services.session_backup import SessionBackupError
+from iac_code.services.session_storage import SessionStorage
 
 logger = logging.getLogger(__name__)
 _V03_JSONRPC_METHODS = frozenset(
@@ -392,6 +395,7 @@ def create_app(
     supported_interfaces: list[dict[str, str]] | None = None,
     agent_extensions: object | None = None,
     auto_approve_permissions: bool = False,
+    permission_wait: object | None = None,
     thinking_exposure: object | None = None,
     idle_shutdown_seconds: float = 0,
     idle_shutdown_callback: Callable[[], None] | None = None,
@@ -423,6 +427,7 @@ def create_app(
         supported_interfaces=supported_interfaces,
         agent_extensions=agent_extensions,
         auto_approve_permissions=auto_approve_permissions,
+        permission_wait=permission_wait,
         thinking_exposure=thinking_exposure,
     )
     from iac_code.a2a.pipeline_recovery import A2APipelineRecoveryService
@@ -475,6 +480,42 @@ def create_app(
     async def get_readiness(request: Request) -> JSONResponse:
         return JSONResponse(configuration_readiness(model=model))
 
+    async def ensure_session_restored(request: Request) -> JSONResponse:
+        try:
+            payload = await request.json()
+        except Exception:
+            return JSONResponse({"error": "Invalid session restore request."}, status_code=400)
+        if not isinstance(payload, dict):
+            return JSONResponse({"error": "Invalid session restore request."}, status_code=400)
+        raw_cwd = payload.get("cwd")
+        raw_session_id = payload.get("sessionId")
+        if not isinstance(raw_cwd, str) or not raw_cwd or not isinstance(raw_session_id, str):
+            return JSONResponse({"error": "Invalid session restore request."}, status_code=400)
+
+        executor = getattr(components.handler, "agent_executor", None)
+        resolve_cwd = getattr(executor, "_resolve_cwd", None)
+        backup_service = components.backup_service
+        if not callable(resolve_cwd) or backup_service is None:
+            return JSONResponse({"error": "Session restore is unavailable."}, status_code=503)
+        try:
+            session_id = validate_protocol_id(raw_session_id)
+            cwd = resolve_cwd({"iac_code": {"cwd": raw_cwd}})
+            result = await asyncio.to_thread(backup_service.reconcile_session, cwd, session_id)
+            exists = SessionStorage().exists(cwd, session_id)
+        except ValueError:
+            return JSONResponse({"error": "Invalid session restore request."}, status_code=400)
+        except SessionBackupError as exc:
+            logger.warning("A2A session restore failed error_type=%s", type(exc).__name__)
+            return JSONResponse({"error": "Unable to restore the A2A session."}, status_code=503)
+        except Exception as exc:
+            logger.warning("A2A session restore failed error_type=%s", type(exc).__name__)
+            return JSONResponse({"error": "Unable to restore the A2A session."}, status_code=503)
+
+        if not exists:
+            return JSONResponse({"status": "not_found"}, status_code=404)
+        status = "restored" if getattr(result, "action", None) == "restored" else "current"
+        return JSONResponse({"status": status})
+
     recovery_service = A2APipelineRecoveryService(task_store=components.task_store)
 
     async def recovery_path_roots(*, context_id: str | None, task_id: str | None) -> list[dict[str, str]]:
@@ -521,6 +562,7 @@ def create_app(
         Route("/health", health, methods=["GET"]),
         Route(AGENT_CARD_WELL_KNOWN_PATH, get_agent_card, methods=["GET"]),
         Route("/iac-code/readiness", get_readiness, methods=["GET"]),
+        Route("/iac-code/session/ensure-restored", ensure_session_restored, methods=["POST"]),
         Route("/iac-code/pipeline/state", get_pipeline_state, methods=["GET"]),
     ]
     install_jsonrpc_error_data_passthrough()
@@ -607,6 +649,7 @@ def run_server(
     push_consumer_name: str | None = None,
     push_lease_timeout_ms: int = 300_000,
     auto_approve_permissions: bool = False,
+    permission_wait: object | None = None,
     thinking_exposure: object | None = None,
     idle_shutdown_seconds: float = 0,
 ) -> None:
@@ -678,6 +721,7 @@ def run_server(
         "push_lease_timeout_ms": push_lease_timeout_ms,
         "supported_interfaces": supported_interfaces,
         "auto_approve_permissions": auto_approve_permissions,
+        "permission_wait": permission_wait,
         "thinking_exposure": thinking_exposure,
     }
 
@@ -783,6 +827,7 @@ def run_server(
             push_lease_timeout_ms=push_lease_timeout_ms,
             supported_interfaces=supported_interfaces,
             auto_approve_permissions=auto_approve_permissions,
+            permission_wait=permission_wait,
             thinking_exposure=thinking_exposure,
             idle_shutdown_seconds=idle_shutdown_seconds,
             idle_shutdown_callback=request_idle_shutdown,

{iac_code-0.13.1 → iac_code-0.14.0}/src/iac_code/a2a/backup.py

@@ -41,13 +41,14 @@ async def backup_session_async(
     critical: bool,
     metrics: Any | None = None,
     publication_proofs: dict[str, BackupPublicationProof] | None = None,
+    backup_call: Callable[..., Any] | None = None,
 ) -> Any | None:
     failed_recorded = False
     try:
         kwargs: dict[str, Any] = {"reason": reason, "critical": critical}
         if publication_proofs is not None:
             kwargs["publication_proofs"] = publication_proofs
-        result = await run_sync_fenced(backup_service.backup_session, cwd, session_id, **kwargs)
+        result = await run_sync_fenced(backup_call or backup_service.backup_session, cwd, session_id, **kwargs)
         retry_count = _retry_count(result)
         if getattr(result, "enabled", False) and not getattr(result, "succeeded", True):
             message = str(

{iac_code-0.13.1 → iac_code-0.14.0}/src/iac_code/a2a/client.py

@@ -154,22 +154,29 @@ class A2AClient:
         *,
         cwd: str,
         context_id: str | None = None,
+        task_id: str | None = None,
+        message_id: str | None = None,
         model: str | None = None,
         iac_code_api_key: str | None = None,
         thinking_enabled: bool | None = None,
         thinking_effort: str | None = None,
         thinking_budget: int | None = None,
+        iac_code_metadata: Mapping[str, Any] | None = None,
     ) -> A2AClientResponse:
         payload = self._message_payload(
             method="SendMessage",
             prompt=prompt,
+            parts=None,
             cwd=cwd,
             context_id=context_id,
+            task_id=task_id,
+            message_id=message_id,
             model=model,
             iac_code_api_key=iac_code_api_key,
             thinking_enabled=thinking_enabled,
             thinking_effort=thinking_effort,
             thinking_budget=thinking_budget,
+            additional_iac_code_metadata=iac_code_metadata,
         )
         transport = self._make_transport_client(url)
         response = await transport.send(payload)
@@ -182,22 +189,101 @@ class A2AClient:
         *,
         cwd: str,
         context_id: str | None = None,
+        task_id: str | None = None,
+        message_id: str | None = None,
         model: str | None = None,
         iac_code_api_key: str | None = None,
         thinking_enabled: bool | None = None,
         thinking_effort: str | None = None,
         thinking_budget: int | None = None,
+        iac_code_metadata: Mapping[str, Any] | None = None,
     ) -> AsyncIterator[dict[str, Any]]:
         payload = self._message_payload(
             method="SendStreamingMessage",
             prompt=prompt,
+            parts=None,
             cwd=cwd,
             context_id=context_id,
+            task_id=task_id,
+            message_id=message_id,
             model=model,
             iac_code_api_key=iac_code_api_key,
             thinking_enabled=thinking_enabled,
             thinking_effort=thinking_effort,
             thinking_budget=thinking_budget,
+            additional_iac_code_metadata=iac_code_metadata,
+        )
+        transport = self._make_transport_client(url)
+        async for event in transport.stream(payload):
+            yield event
+
+    async def send_message_parts(
+        self,
+        url: str,
+        parts: Sequence[Mapping[str, Any]],
+        *,
+        cwd: str,
+        context_id: str,
+        task_id: str | None = None,
+        message_id: str | None = None,
+        model: str | None = None,
+        iac_code_api_key: str | None = None,
+        thinking_enabled: bool | None = None,
+        thinking_effort: str | None = None,
+        thinking_budget: int | None = None,
+        iac_code_metadata: Mapping[str, Any] | None = None,
+    ) -> A2AClientResponse:
+        """Send pre-validated A2A parts while preserving iac-code request metadata."""
+        payload = self._message_payload(
+            method="SendMessage",
+            prompt=None,
+            parts=parts,
+            cwd=cwd,
+            context_id=context_id,
+            task_id=task_id,
+            message_id=message_id,
+            model=model,
+            iac_code_api_key=iac_code_api_key,
+            thinking_enabled=thinking_enabled,
+            thinking_effort=thinking_effort,
+            thinking_budget=thinking_budget,
+            additional_iac_code_metadata=iac_code_metadata,
+        )
+        transport = self._make_transport_client(url)
+        response = await transport.send(payload)
+        return A2AClientResponse(payload=response)
+
+    async def stream_message_parts(
+        self,
+        url: str,
+        parts: Sequence[Mapping[str, Any]],
+        *,
+        cwd: str,
+        context_id: str,
+        task_id: str | None = None,
+        message_id: str | None = None,
+        model: str | None = None,
+        iac_code_api_key: str | None = None,
+        thinking_enabled: bool | None = None,
+        thinking_effort: str | None = None,
+        thinking_budget: int | None = None,
+        iac_code_metadata: Mapping[str, Any] | None = None,
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Stream a message whose parts were validated by an embedding protocol adapter."""
+        payload = self._message_payload(
+            method="SendStreamingMessage",
+            prompt=None,
+            parts=parts,
+            cwd=cwd,
+            context_id=context_id,
+            task_id=task_id,
+            message_id=message_id,
+            model=model,
+            iac_code_api_key=iac_code_api_key,
+            thinking_enabled=thinking_enabled,
+            thinking_effort=thinking_effort,
+            thinking_budget=thinking_budget,
+            additional_iac_code_metadata=iac_code_metadata,
         )
         transport = self._make_transport_client(url)
         async for event in transport.stream(payload):
@@ -209,6 +295,43 @@ class A2AClient:
             params["historyLength"] = history_length
         return await self._send_jsonrpc(url, method="GetTask", params=params)
 
+    async def get_pipeline_state(
+        self,
+        url: str,
+        *,
+        task_id: str,
+        after_sequence: int | None = None,
+    ) -> dict[str, Any] | None:
+        """Read iac-code's existing A2A pipeline recovery extension over HTTP."""
+
+        endpoint = url.rstrip("/") + "/iac-code/pipeline/state"
+        params = _without_none({"taskId": task_id, "afterSequence": after_sequence})
+        response = await self._http_client.get(endpoint, params=params, headers=self._jsonrpc_headers())
+        if response.status_code == 404:
+            return None
+        response.raise_for_status()
+        data = response.json()
+        if not isinstance(data, dict):
+            raise ValueError("A2A pipeline state response must be a JSON object")
+        return data
+
+    async def ensure_session_restored(self, url: str, *, cwd: str, session_id: str) -> bool:
+        """Ensure the A2A session payload is locally available before a cold resume."""
+
+        endpoint = url.rstrip("/") + "/iac-code/session/ensure-restored"
+        response = await self._http_client.post(
+            endpoint,
+            json={"cwd": cwd, "sessionId": session_id},
+            headers=self._jsonrpc_headers(),
+        )
+        if response.status_code == 404:
+            return False
+        response.raise_for_status()
+        data = response.json()
+        if not isinstance(data, dict) or data.get("status") not in {"current", "restored"}:
+            raise ValueError("A2A session restore response must report a ready session")
+        return True
+
     async def list_tasks(
         self,
         url: str,
@@ -397,16 +520,23 @@ class A2AClient:
         self,
         *,
         method: str,
-        prompt: str,
+        prompt: str | None,
+        parts: Sequence[Mapping[str, Any]] | None,
         cwd: str,
         context_id: str | None,
+        task_id: str | None,
+        message_id: str | None,
         model: str | None,
         iac_code_api_key: str | None,
         thinking_enabled: bool | None,
         thinking_effort: str | None,
         thinking_budget: int | None,
+        additional_iac_code_metadata: Mapping[str, Any] | None,
     ) -> dict[str, Any]:
-        iac_code_metadata: dict[str, Any] = {"cwd": cwd}
+        if (prompt is None) == (parts is None):
+            raise ValueError("Exactly one of prompt or parts is required")
+        iac_code_metadata: dict[str, Any] = dict(additional_iac_code_metadata or {})
+        iac_code_metadata["cwd"] = cwd
         if model:
             stripped_model = model.strip()
             if stripped_model:
@@ -427,13 +557,15 @@ class A2AClient:
         if thinking:
             iac_code_metadata["thinking"] = thinking
         message: dict[str, Any] = {
-            "messageId": str(uuid.uuid4()),
+            "messageId": message_id or str(uuid.uuid4()),
             "role": "ROLE_USER",
-            "parts": [{"text": prompt}],
+            "parts": [{"text": prompt}] if prompt is not None else [dict(part) for part in parts or ()],
             "metadata": {"iac_code": iac_code_metadata},
         }
         if context_id:
             message["contextId"] = context_id
+        if task_id:
+            message["taskId"] = task_id
         return self._jsonrpc_payload(
             method=method,
             params={

{iac_code-0.13.1 → iac_code-0.14.0}/src/iac_code/a2a/events.py

@@ -82,7 +82,13 @@ def _truncate(value: Any, *, _depth: int = 0) -> Any:
 
 
 def _public_tool_input_metadata(tool_name: str, tool_input: dict[str, Any]) -> dict[str, Any]:
-    return {"inputSummary": build_input_summary(tool_name, tool_input)}
+    return {
+        "inputSummary": build_input_summary(tool_name, tool_input),
+        # This is protocol data consumed as AG-UI TOOL_CALL_ARGS, not an audit
+        # record. Keep the canonical arguments here; the A2A wire boundary
+        # applies the existing path-only safe-mode projection to a copy.
+        "toolInput": copy.deepcopy(tool_input),
+    }
 
 
 def _public_permission_input_metadata(tool_name: str, tool_input: dict[str, Any]) -> dict[str, Any]:
@@ -288,6 +294,9 @@ async def publish_stream_event(
     auto_approve_permissions: bool = False,
     exposure_types: Any = None,
     iac_code_session_id: str | None = None,
+    permission_wait_cwd: str | None = None,
+    permission_wait_backup_service: Any | None = None,
+    permission_wait_metrics: Any | None = None,
 ) -> str | None:
     enabled_exposure_types = normalize_a2a_exposure_types(exposure_types)
 
@@ -437,56 +446,18 @@ async def publish_stream_event(
     permission_event = _permission_request_event(event)
     if permission_event is not None:
         if permission_input_registry is not None and permission_resolver is None and not auto_approve_permissions:
-            pending = await permission_input_registry.register(
-                permission_event,
+            await publish_interactive_permission_boundary(
+                event_queue,
+                permission_event=permission_event,
+                permission_input_registry=permission_input_registry,
                 task_id=task_id,
                 context_id=context_id,
+                iac_code_session_id=iac_code_session_id,
+                permission_wait_cwd=permission_wait_cwd,
+                permission_wait_backup_service=permission_wait_backup_service,
+                permission_wait_metrics=permission_wait_metrics,
+                wait_for_response=True,
             )
-            try:
-                await _enqueue_status(
-                    event_queue,
-                    task_id=task_id,
-                    context_id=context_id,
-                    state=TaskState.TASK_STATE_INPUT_REQUIRED,
-                    metadata={
-                        "iac_code": {
-                            "input": pending.envelope(),
-                            "permission": {
-                                "autoApproved": False,
-                                "pending": True,
-                                "toolName": permission_event.tool_name,
-                                "toolUseId": permission_event.tool_use_id,
-                            },
-                        }
-                    },
-                    iac_code_session_id=iac_code_session_id,
-                )
-                future = permission_event.response_future
-                if future is None:
-                    await permission_input_registry.fail(pending)
-                    return None
-                await asyncio.shield(future)
-                await _enqueue_status(
-                    event_queue,
-                    task_id=task_id,
-                    context_id=context_id,
-                    state=TaskState.TASK_STATE_WORKING,
-                    metadata={
-                        "iac_code": {
-                            "inputReceived": {
-                                "kind": "permission",
-                                "inputId": pending.input_id,
-                                "toolUseId": permission_event.tool_use_id,
-                            }
-                        }
-                    },
-                    iac_code_session_id=iac_code_session_id,
-                )
-            except BaseException:
-                await permission_input_registry.fail(pending)
-                raise
-            finally:
-                await permission_input_registry.complete(pending)
             return None
         approved = auto_approve_permissions
         if permission_resolver is not None:
@@ -525,6 +496,20 @@ async def publish_stream_event(
         return None
 
     if isinstance(event, MessageEndEvent):
+        usage = {
+            "inputTokens": event.usage.input_tokens,
+            "outputTokens": event.usage.output_tokens,
+            "totalTokens": event.usage.total_tokens,
+        }
+        provider = getattr(event.usage, "provider", None)
+        model = getattr(event.usage, "model", None)
+        cached_input_tokens = getattr(event.usage, "cache_read_input_tokens", None)
+        if isinstance(provider, str) and provider:
+            usage["provider"] = provider
+        if isinstance(model, str) and model:
+            usage["model"] = model
+        if isinstance(cached_input_tokens, int) and not isinstance(cached_input_tokens, bool):
+            usage["cachedInputTokens"] = cached_input_tokens
         await _enqueue_status(
             event_queue,
             task_id=task_id,
@@ -532,11 +517,7 @@ async def publish_stream_event(
             state=TaskState.TASK_STATE_WORKING,
             metadata={
                 "iac_code": {
-                    "usage": {
-                        "inputTokens": event.usage.input_tokens,
-                        "outputTokens": event.usage.output_tokens,
-                        "totalTokens": event.usage.total_tokens,
-                    }
+                    "usage": usage
                 }
             },
             iac_code_session_id=iac_code_session_id,
@@ -571,6 +552,107 @@ async def publish_stream_event(
     return None
 
 
+async def publish_interactive_permission_boundary(
+    event_queue: Any,
+    *,
+    permission_event: PermissionRequestEvent,
+    permission_input_registry: Any,
+    task_id: str,
+    context_id: str,
+    iac_code_session_id: str | None,
+    permission_wait_cwd: str | None,
+    permission_wait_backup_service: Any | None,
+    permission_wait_metrics: Any | None = None,
+    wait_for_response: bool,
+) -> Any:
+    """Publish one real external permission wait, optionally detaching Normal SSE."""
+
+    pending = await permission_input_registry.register(
+        permission_event,
+        task_id=task_id,
+        context_id=context_id,
+    )
+    try:
+        if (
+            iac_code_session_id is not None
+            and permission_wait_cwd is not None
+            and permission_wait_backup_service is not None
+        ):
+            await permission_input_registry.open_durable_boundary(
+                pending,
+                cwd=permission_wait_cwd,
+                session_id=iac_code_session_id,
+                permission_class="normal",
+                backup_service=permission_wait_backup_service,
+                metrics=permission_wait_metrics,
+            )
+        await _enqueue_status(
+            event_queue,
+            task_id=task_id,
+            context_id=context_id,
+            state=TaskState.TASK_STATE_INPUT_REQUIRED,
+            metadata={
+                "iac_code": {
+                    "input": pending.envelope(),
+                    "permission": {
+                        "autoApproved": False,
+                        "pending": True,
+                        "toolName": permission_event.tool_name,
+                        "toolUseId": permission_event.tool_use_id,
+                    },
+                }
+            },
+            iac_code_session_id=iac_code_session_id,
+        )
+        if not wait_for_response:
+            return pending
+        future = permission_event.response_future
+        if future is None:
+            await permission_input_registry.fail(pending)
+            return pending
+        outcome = await asyncio.shield(future)
+        from iac_code.types.stream_events import PermissionWaitOutcome
+
+        if outcome is PermissionWaitOutcome.SUSPEND:
+            return pending
+        await publish_permission_input_received(
+            event_queue,
+            pending=pending,
+            iac_code_session_id=iac_code_session_id,
+        )
+        return pending
+    except BaseException:
+        await permission_input_registry.fail(pending)
+        raise
+    finally:
+        if wait_for_response:
+            await permission_input_registry.complete(pending)
+
+
+async def publish_permission_input_received(
+    event_queue: Any,
+    *,
+    pending: Any,
+    iac_code_session_id: str | None,
+) -> None:
+    await _enqueue_status(
+        event_queue,
+        task_id=pending.task_id,
+        context_id=pending.context_id,
+        state=TaskState.TASK_STATE_WORKING,
+        metadata={
+            "iac_code": {
+                "inputReceived": {
+                    "kind": "permission",
+                    "inputId": pending.input_id,
+                    "toolUseId": pending.request.tool_use_id,
+                }
+            }
+        },
+        iac_code_session_id=iac_code_session_id,
+    )
+
+
 def _emit_auto_permission_audit(
     request: PermissionRequestEvent,
     approved: bool,
@@ -601,6 +683,8 @@ def _emit_resolver_permission_audit(
     *,
     persistence_failure: bool = False,
 ) -> bool:
+    if request.permission_decision_audited and not persistence_failure:
+        return True
     source = "a2a_resolver"
     reason_type = "a2a_resolver"
     reason_detail = "allow" if approved else "deny"