machinaos 0.0.76 → 0.0.78

package/server/services/cli_agent/types.py

@@ -0,0 +1,227 @@
+"""Pydantic task specs (discriminated union) + result models.
+
+The discriminated union lets each plugin's `Params.tasks` be hard-typed
+to one variant (e.g. `list[ClaudeTaskSpec]`), so the LLM tool schema
+fast-path at `services/ai.py:2898` produces a clean per-provider schema
+with no `$defs`/`$ref`.
+
+v1 ships Claude + Codex plugins fully wired. `GeminiTaskSpec` is part of
+the union (and exported) so the type system, factory dispatch keys, and
+config JSON are ready, but `create_cli_provider("gemini")` raises
+`NotImplementedError`. v2 swaps the stub for the real impl.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from services.cli_agent.protocol import CanonicalUsage
+
+
+# ---------------------------------------------------------------------------
+# Task specs (discriminated union)
+# ---------------------------------------------------------------------------
+
+class BaseAICliTaskSpec(BaseModel):
+    """Shared fields for every CLI task."""
+    task_id: Optional[str] = Field(
+        default=None,
+        description="Auto-assigned `t_<8hex>` if omitted.",
+    )
+    prompt: str = Field(
+        ...,
+        description="The task prompt sent to the CLI.",
+        json_schema_extra={"rows": 4},
+    )
+    branch: Optional[str] = Field(
+        default=None,
+        description="Branch name for the per-task git worktree. "
+                    "Auto-named `machina/<task_id>` if omitted.",
+    )
+    model: Optional[str] = Field(default=None)
+    timeout_seconds: int = Field(
+        default=600, ge=10, le=3600,
+        description="Hard timeout per task. On expiry the session is "
+                    "terminate_then_kill'd and a diagnostic dump is "
+                    "written to ~/.claude-machina/logs/.",
+    )
+    system_prompt: Optional[str] = Field(
+        default=None,
+        json_schema_extra={"rows": 3},
+    )
+
+    # `extra="forbid"` surfaces typo'd task fields as Pydantic
+    # ValidationError at the spec boundary instead of silently dropping
+    # them and confusing downstream consumers.
+    model_config = ConfigDict(extra="forbid")
+
+
+class ClaudeTaskSpec(BaseAICliTaskSpec):
+    """Claude Code CLI task. Full feature set: sessions, resume, budget,
+    turns, allowed_tools, permission_mode."""
+    provider: Literal["claude"] = "claude"
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Start a named session. Pair with `resume_session_id` "
+                    "to chain conversations.",
+    )
+    resume_session_id: Optional[str] = Field(
+        default=None,
+        description="Resume from a prior session_id (returned in a previous "
+                    "task's `result` event).",
+    )
+    max_turns: Optional[int] = Field(
+        default=None, ge=1,
+        description="Per-task turn cap. Defaults to provider config.",
+    )
+    max_budget_usd: Optional[float] = Field(
+        default=None, ge=0,
+        description="Per-task USD budget. Defaults to provider config.",
+    )
+    allowed_tools: Optional[str] = Field(
+        default=None,
+        description="Comma-separated tool list. Defaults to "
+                    "Read,Edit,Bash,Glob,Grep,Write.",
+    )
+    permission_mode: Literal["default", "acceptEdits", "plan", "auto", "dontAsk", "bypassPermissions"] = "acceptEdits"
+
+    # ---- optional documented CLI flags (cli-reference) ----
+    effort: Optional[Literal["low", "medium", "high", "xhigh", "max"]] = Field(
+        default=None,
+        description="Reasoning-effort level. Available levels depend on the model.",
+    )
+    fallback_model: Optional[str] = Field(
+        default=None,
+        description="Fallback model when the primary is overloaded (print mode).",
+    )
+    add_dir: List[str] = Field(
+        default_factory=list,
+        description="Extra working directories the CLI may read/edit.",
+    )
+    disallowed_tools: Optional[str] = Field(
+        default=None,
+        description="Comma-separated tools removed from Claude's context.",
+    )
+    agent: Optional[str] = Field(
+        default=None,
+        description="Override the configured `agent` setting for this task.",
+    )
+
+
+class CodexTaskSpec(BaseAICliTaskSpec):
+    """OpenAI Codex CLI task. Sandbox-first; no session/resume/budget/turns."""
+    provider: Literal["codex"] = "codex"
+    sandbox: Literal["read-only", "workspace-write", "danger-full-access"] = (
+        "workspace-write"
+    )
+    ask_for_approval: Literal["untrusted", "on-request", "never"] = "never"
+
+
+class GeminiTaskSpec(BaseAICliTaskSpec):
+    """Google Gemini CLI task (v2 — stub provider in v1).
+
+    Schema lives in v1 so the discriminated union JSON Schema for the
+    LLM tool fast-path doesn't change when v2 lands.
+    """
+    provider: Literal["gemini"] = "gemini"
+    session_id: Optional[str] = None
+    resume: Optional[str] = Field(
+        default=None,
+        description='"latest" | "<index>" | "<UUID>"',
+    )
+    yolo: bool = False
+    sandbox: bool = False
+
+
+AICliTaskSpec = Annotated[
+    Union[ClaudeTaskSpec, CodexTaskSpec, GeminiTaskSpec],
+    Field(discriminator="provider"),
+]
+
+
+# ---------------------------------------------------------------------------
+# Result models (Pydantic for serialisation; mirror dataclasses in protocol.py)
+# ---------------------------------------------------------------------------
+
+class CanonicalUsagePydantic(BaseModel):
+    """Pydantic mirror of `protocol.CanonicalUsage` for output serialisation."""
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read: int = 0
+    cache_write: int = 0
+    reasoning_tokens: int = 0
+    request_count: int = 0
+
+
+class SessionResultModel(BaseModel):
+    """Per-task result, JSON-serialisable."""
+    task_id: str
+    session_id: Optional[str] = None
+    provider: str = ""
+    prompt: str = ""
+    branch: Optional[str] = None
+    worktree_path: Optional[str] = None
+    response: str = ""
+    cost_usd: Optional[float] = None
+    duration_ms: Optional[int] = None
+    num_turns: Optional[int] = None
+    tool_calls: int = 0
+    canonical_usage: CanonicalUsagePydantic = Field(
+        default_factory=CanonicalUsagePydantic,
+    )
+    provider_data: Dict[str, Any] = Field(default_factory=dict)
+    success: bool = False
+    error: Optional[str] = None
+
+
+class BatchSummary(BaseModel):
+    """Aggregated batch summary."""
+    n_tasks: int = 0
+    n_succeeded: int = 0
+    n_failed: int = 0
+    total_cost_usd: Optional[float] = None
+    wall_clock_ms: int = 0
+    budget_remaining_usd: Optional[float] = None
+
+
+class BatchResultModel(BaseModel):
+    """Top-level batch result returned by `run_batch()`."""
+    tasks: List[SessionResultModel] = Field(default_factory=list)
+    summary: BatchSummary = Field(default_factory=BatchSummary)
+    provider: str = ""
+    timestamp: str = ""
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def session_result_to_model(sr: Any) -> SessionResultModel:
+    """Convert a `protocol.SessionResult` dataclass to its Pydantic mirror."""
+    cu = CanonicalUsagePydantic(
+        input_tokens=sr.canonical_usage.input_tokens,
+        output_tokens=sr.canonical_usage.output_tokens,
+        cache_read=sr.canonical_usage.cache_read,
+        cache_write=sr.canonical_usage.cache_write,
+        reasoning_tokens=sr.canonical_usage.reasoning_tokens,
+        request_count=sr.canonical_usage.request_count,
+    )
+    return SessionResultModel(
+        task_id=sr.task_id,
+        session_id=sr.session_id,
+        provider=sr.provider,
+        prompt=sr.prompt,
+        branch=sr.branch,
+        worktree_path=sr.worktree_path,
+        response=sr.response,
+        cost_usd=sr.cost_usd,
+        duration_ms=sr.duration_ms,
+        num_turns=sr.num_turns,
+        tool_calls=sr.tool_calls,
+        canonical_usage=cu,
+        provider_data=sr.provider_data,
+        success=sr.success,
+        error=sr.error,
+    )

package/server/services/cli_agent/workflow_tools.py

@@ -0,0 +1,233 @@
+"""Per-batch workflow-tool bridge for the FastMCP server.
+
+When an agent batch wires nodes through ``input-tools``, each connected
+node is exposed on the FastMCP server as its own
+``mcp__machinaos__<node_type>`` entry. The spawned ``claude -p`` sees
+those tools on the very first ``tools/list`` and can invoke them
+directly — no two-step generic-wrapper indirection.
+
+Schema generation is delegated to FastMCP: we build a function whose
+``inspect.signature`` mirrors the plugin's Pydantic ``Params`` field-
+for-field, so FastMCP advertises a flat ``inputSchema`` matching the
+plugin Params. Per-batch scoping is enforced inside the handler via
+``_require_batch()`` so concurrent batches sharing the same tool name
+are isolated; refcounts (``add_tool`` on first wire, ``remove_tool`` on
+last unwire) keep the FastMCP registry tidy.
+
+Public API:
+  - :func:`expose_workflow_tools(connected_tools)` — call from
+    ``mcp_server.register_batch``
+  - :func:`unexpose_workflow_tools(connected_tools)` — call from
+    ``mcp_server.unregister_batch``
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from typing import Any, Dict, List, Optional
+
+from core.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# node_type -> count of active batches that wired it. Tools are added
+# to FastMCP on first wire, removed on last unwire — per-handler scope
+# checks ``_require_batch`` against the calling batch's connected_tools.
+_active_tool_refcounts: Dict[str, int] = {}
+
+
+def expose_workflow_tools(connected_tools: List[Dict[str, Any]]) -> None:
+    """Add one MCP tool per connected workflow-tool node."""
+    mcp = _get_mcp()
+    if mcp is None or not connected_tools:
+        return
+    from services.node_registry import get_node_class
+
+    for entry in connected_tools:
+        node_type = entry.get("node_type")
+        if not node_type:
+            continue
+        prev = _active_tool_refcounts.get(node_type, 0)
+        _active_tool_refcounts[node_type] = prev + 1
+        if prev > 0:
+            continue  # already exposed by another concurrent batch
+        cls = get_node_class(node_type)
+        if cls is None or getattr(cls, "Params", None) is None:
+            logger.warning(
+                "[CC-Agent MCP] cannot expose %s: class or Params missing",
+                node_type,
+            )
+            continue
+        try:
+            handler = _build_handler(node_type, cls.Params)
+            mcp.add_tool(
+                handler, name=node_type,
+                description=(getattr(cls, "description", None)
+                             or f"MachinaOs workflow tool: {node_type}"),
+            )
+            logger.info("[CC-Agent MCP] exposed mcp__machinaos__%s", node_type)
+        except Exception as exc:  # pragma: no cover
+            logger.warning("[CC-Agent MCP] add_tool(%s) failed: %s", node_type, exc)
+    _schedule_list_changed_notify()
+
+
+def unexpose_workflow_tools(connected_tools: List[Dict[str, Any]]) -> None:
+    """Decrement refcount; remove the MCP tool when no batch references it."""
+    mcp = _get_mcp()
+    if mcp is None:
+        return
+    for entry in connected_tools:
+        node_type = entry.get("node_type")
+        if not node_type:
+            continue
+        remaining = _active_tool_refcounts.get(node_type, 0) - 1
+        if remaining > 0:
+            _active_tool_refcounts[node_type] = remaining
+            continue
+        _active_tool_refcounts.pop(node_type, None)
+        try:
+            mcp.remove_tool(node_type)
+            logger.info("[CC-Agent MCP] removed mcp__machinaos__%s", node_type)
+        except Exception as exc:  # pragma: no cover
+            logger.debug("[CC-Agent MCP] remove_tool(%s): %s", node_type, exc)
+    _schedule_list_changed_notify()
+
+
+# ---------------------------------------------------------------------------
+# Internals
+# ---------------------------------------------------------------------------
+
+def _get_mcp() -> Optional[Any]:
+    """Late import to avoid a circular import with ``mcp_server``."""
+    from services.cli_agent import mcp_server
+    return mcp_server._mcp_singleton
+
+
+def _schedule_list_changed_notify() -> None:
+    """Fire-and-forget ``notifications/tools/list_changed`` to the
+    connected MCP client.
+
+    FastMCP does NOT emit this automatically on ``add_tool`` /
+    ``remove_tool`` (verified at
+    ``mcp/server/fastmcp/tools/tool_manager.py`` — both methods only
+    mutate the ``_tools`` dict, with no notification dispatch). Without
+    this manual notify, tools registered after the agent's first
+    ``tools/list`` request stay invisible until reconnect.
+
+    Best-effort: requires a running asyncio loop (always true in the
+    ``service.run_batch`` call path). Failures log at WARN and don't
+    abort the batch.
+    """
+    mcp = _get_mcp()
+    if mcp is None:
+        return
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return  # no running loop — sync test path
+
+    async def _do_notify() -> None:
+        try:
+            session = (
+                getattr(mcp, "session", None)
+                or getattr(getattr(mcp, "_mcp_server", None), "session", None)
+            )
+            if session is not None and hasattr(session, "send_tool_list_changed"):
+                await session.send_tool_list_changed()
+        except Exception as exc:
+            logger.warning(
+                "[CC-Agent workflow_tools] tools/list_changed notify failed: %s",
+                exc,
+            )
+
+    loop.create_task(_do_notify())
+
+
+def _build_handler(node_type: str, params_cls: type):
+    """Build an async handler whose ``inspect.signature`` mirrors the
+    plugin Params field-for-field. FastMCP iterates that signature to
+    derive the ``inputSchema``; flat fields → flat MCP arguments.
+    """
+    from pydantic_core import PydanticUndefined
+
+    fields = getattr(params_cls, "model_fields", {}) or {}
+    parameters: List[inspect.Parameter] = []
+    annotations: Dict[str, Any] = {"return": Dict[str, Any]}
+    for fname, finfo in fields.items():
+        annotations[fname] = finfo.annotation
+        if finfo.is_required():
+            default: Any = inspect.Parameter.empty
+        elif finfo.default_factory is not None:  # type: ignore[truthy-function]
+            try:
+                default = finfo.default_factory()  # materialize once for the signature
+            except Exception:
+                default = None
+        elif finfo.default is PydanticUndefined:
+            default = None
+        else:
+            default = finfo.default
+        parameters.append(inspect.Parameter(
+            fname, inspect.Parameter.KEYWORD_ONLY,
+            annotation=finfo.annotation, default=default,
+        ))
+
+    async def _handler(**kwargs: Any) -> Dict[str, Any]:
+        from services.cli_agent.mcp_server import _require_batch
+        ctx = _require_batch()
+        entry = next(
+            (t for t in ctx.connected_tools if t.get("node_type") == node_type),
+            None,
+        )
+        if entry is None:
+            return {
+                "error": f"tool {node_type!r} not connected to this batch",
+                "status": 403,
+            }
+        # Validate via the plugin's own Params model, then dump back to
+        # a plain dict for execute_tool. `exclude_unset=True` drops
+        # defaults the agent didn't supply.
+        try:
+            validated = params_cls(**kwargs)
+            args = validated.model_dump(exclude_unset=True)
+        except Exception:
+            args = dict(kwargs)
+        logger.info(
+            "[CC-Agent MCP %s] node=%s wf=%s args_keys=%s",
+            node_type, ctx.node_id, ctx.workflow_id, list(args.keys()),
+        )
+        from services.handlers.tools import execute_tool
+        config: Dict[str, Any] = {
+            "node_type": node_type,
+            "node_id": entry.get("node_id"),
+            "workflow_id": ctx.workflow_id,
+            "label": entry.get("label") or node_type,
+            "parameters": dict(entry.get("parameters") or {}),
+        }
+        result = await execute_tool(node_type, args, config)
+        if not isinstance(result, dict):
+            result = {"result": result}
+        if "error" in result:
+            logger.warning(
+                "[CC-Agent MCP %s] node=%s ERROR: %s",
+                node_type, entry.get("node_id"), result.get("error"),
+            )
+        else:
+            logger.info(
+                "[CC-Agent MCP %s] node=%s OK (result_keys=%s)",
+                node_type, entry.get("node_id"), list(result.keys())[:8],
+            )
+        return result
+
+    _handler.__name__ = node_type
+    _handler.__annotations__ = annotations
+    _handler.__signature__ = inspect.Signature(  # type: ignore[attr-defined]
+        parameters, return_annotation=Dict[str, Any],
+    )
+    return _handler
+
+
+def _reset_for_tests() -> None:  # pragma: no cover
+    """Wipe the refcount registry. ONLY use in tests."""
+    _active_tool_refcounts.clear()

package/server/services/credential_registry.py

@@ -44,6 +44,10 @@ class CredentialRegistry:
         self._raw: Optional[Dict[str, Any]] = None
         self._resolved_providers: Optional[Dict[str, Dict[str, Any]]] = None
         self._version: Optional[str] = None
+        # Bumped on every credential mutation; folded into the version
+        # hash so the conditional fetch returns fresh ``stored`` flags.
+        # See ``get_version()`` + ``invalidate_version()``.
+        self._mutation_seq: int = 0
 
     @classmethod
     def get_instance(cls) -> "CredentialRegistry":
@@ -162,21 +166,42 @@ class CredentialRegistry:
         return out
 
     def get_version(self) -> str:
-        """Content-sha256 of the resolved catalogue (providers + categories).
+        """Content-sha256 of the resolved catalogue (providers + categories +
+        mutation counter).
 
         Used by clients for warm-start cache invalidation: when the hash
         changes, the client fetches a fresh catalogue; otherwise it serves
         from IndexedDB with zero network traffic.
+
+        The mutation counter is bumped by ``invalidate_version()`` on every
+        credential save / delete / oauth-disconnect — this is what makes
+        the conditional fetch ("`since: <prior version>`") actually return
+        fresh data after a mutation. Without it the hash would only depend
+        on the static catalogue JSON, the version would be constant for
+        the life of the process, and the conditional fetch would always
+        return ``{unchanged: true}`` even after a key was deleted — leaving
+        the per-provider ``stored`` flag stuck at ``true`` on every client.
         """
         if self._version is None:
             payload = {
                 "providers": self.get_all_providers(),
                 "categories": self.get_categories(),
+                "mutation_seq": self._mutation_seq,
             }
             encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
             self._version = hashlib.sha256(encoded).hexdigest()
         return self._version
 
+    def invalidate_version(self) -> None:
+        """Bump the mutation counter so the next ``get_version()`` returns a
+        new hash. Call this from every credential mutation path (store /
+        remove / oauth-disconnect) so the frontend's conditional fetch
+        actually returns fresh ``stored`` flags. Cheap — just a counter
+        increment + clearing the cached hash.
+        """
+        self._mutation_seq += 1
+        self._version = None
+
     def get_catalogue(self) -> Dict[str, Any]:
         """Full payload returned by the `get_credential_catalogue` WebSocket handler."""
         return {