Sanitize generated pricing fields

agent_builder/tools.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import asyncio
+import ast
 import base64
 import io
 import json
@@ -195,6 +196,7 @@ _AGENT_NAME_RE = re.compile(r"^[a-z][a-z0-9-]{1,62}$")
 _SKILL_NAME_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,62}[a-z0-9])?$")
 _BUILDER_STATE_FILE = ".a2a-builder-state.json"
 _BUILDER_INTERNAL_PREFIX = ".agent-builder/"
+_DERIVED_PRICING_FIELDS = frozenset({"compute", "total_usd"})
 
 
 def _validate_name(name: str) -> None:
@@ -307,9 +309,21 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         rel = path.lstrip("/")
         if rel == _BUILDER_STATE_FILE or rel.startswith(_BUILDER_INTERNAL_PREFIX):
             return json.dumps({"error": f"{rel} is managed by agent-builder"})
+        sanitized_fields: tuple[str, ...] = ()
+        if rel == "agent.py":
+            content, sanitized_fields = _strip_unsupported_pricing_fields(content)
         key = prefix + rel
-        store.put(key, content.encode("utf-8"), "text/plain; charset=utf-8")
-        return json.dumps({"ok": True, "path": key, "size": len(content)})
+        encoded = content.encode("utf-8")
+        out: dict[str, Any] = {"ok": True, "path": key, "size": len(encoded)}
+        if sanitized_fields:
+            out["sanitized_agent_card_fields"] = list(sanitized_fields)
+            out["warning"] = (
+                "Removed derived Pricing fields that the Agent Card schema "
+                "does not accept. Declare Resources separately; compute and "
+                "total_usd are filled by the platform."
+            )
+        store.put(key, encoded, "text/plain; charset=utf-8")
+        return json.dumps(out)
 
     @tool
     def read_agent_file(name: str, path: str) -> str:
@@ -388,6 +402,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             prefix = _agent_prefix(name)
         except ValueError as exc:
             return json.dumps({"error": str(exc)})
+        sanitized_fields = _sanitize_agent_workspace_pricing(store, prefix)
 
         # Bundle the workspace files into the sandbox via base64 tar so
         # the sandbox doesn't need MinIO read access for this skill.
@@ -428,11 +443,14 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         if r.status_code >= 400:
             return json.dumps({"error": f"sandbox {r.status_code}", "detail": r.text[:1000]})
         out = r.json()
-        return json.dumps({
+        result: dict[str, Any] = {
             "exit_code": out.get("exit_code"),
             "stdout": (out.get("stdout") or "")[:4000],
             "stderr": (out.get("stderr") or "")[:4000],
-        })
+        }
+        if sanitized_fields:
+            result["sanitized_agent_card_fields"] = list(sanitized_fields)
+        return json.dumps(result)
 
     @tool
     async def cp_deploy_tarball(
@@ -492,6 +510,7 @@ def build_tools(ctx: ToolContext) -> list[Any]:
         if force:
             base_head = None
 
+        sanitized_fields = _sanitize_agent_workspace_pricing(store, prefix)
         bundle_bytes = _tarball_workspace_dir(store, prefix)
         if not bundle_bytes:
             return json.dumps({"error": f"no files at agents/{name}/"})
@@ -561,6 +580,8 @@ def build_tools(ctx: ToolContext) -> list[Any]:
             "workspace_base_head_sha": base_head,
             "live": live,
         }
+        if sanitized_fields:
+            out["sanitized_agent_card_fields"] = list(sanitized_fields)
         if live:
             out["live_version"] = live_card.get("version")
             out["live_skills"] = [
@@ -897,6 +918,109 @@ def _tarball_workspace_dir(
     return buf.getvalue() if added else b""
 
 
+def _sanitize_agent_workspace_pricing(
+    store: _ObjectStore,
+    prefix: str,
+) -> tuple[str, ...]:
+    key = prefix + "agent.py"
+    try:
+        source = store.get(key).decode("utf-8", errors="replace")
+    except Exception:  # noqa: BLE001
+        return ()
+    sanitized, fields = _strip_unsupported_pricing_fields(source)
+    if fields and sanitized != source:
+        store.put(key, sanitized.encode("utf-8"), "text/plain; charset=utf-8")
+    return fields
+
+
+def _strip_unsupported_pricing_fields(source: str) -> tuple[str, tuple[str, ...]]:
+    """Remove platform-derived pricing kwargs from generated Agent source."""
+    try:
+        tree = ast.parse(source)
+    except SyntaxError:
+        return source, ()
+
+    offsets = _line_offsets(source)
+    ranges: list[tuple[int, int, str]] = []
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call) or not _is_pricing_call(node.func):
+            continue
+        for keyword in node.keywords:
+            if keyword.arg not in _DERIVED_PRICING_FIELDS:
+                continue
+            if (
+                keyword.lineno is None
+                or keyword.end_lineno is None
+                or keyword.col_offset is None
+                or keyword.end_col_offset is None
+            ):
+                continue
+            start = offsets[keyword.lineno - 1] + keyword.col_offset
+            end = offsets[keyword.end_lineno - 1] + keyword.end_col_offset
+            ranges.append((start, end, keyword.arg))
+
+    if not ranges:
+        return source, ()
+
+    sanitized = source
+    removed: set[str] = set()
+    for start, end, field in sorted(ranges, key=lambda item: item[0], reverse=True):
+        sanitized = _remove_keyword_source_range(sanitized, start, end)
+        removed.add(field)
+    return sanitized, tuple(sorted(removed))
+
+
+def _line_offsets(source: str) -> list[int]:
+    offsets: list[int] = []
+    offset = 0
+    for line in source.splitlines(keepends=True):
+        offsets.append(offset)
+        offset += len(line)
+    return offsets or [0]
+
+
+def _is_pricing_call(func: ast.expr) -> bool:
+    if isinstance(func, ast.Name):
+        return func.id == "Pricing"
+    if isinstance(func, ast.Attribute):
+        return func.attr == "Pricing"
+    return False
+
+
+def _remove_keyword_source_range(source: str, start: int, end: int) -> str:
+    line_start = source.rfind("\n", 0, start) + 1
+    if source[line_start:start].strip() == "":
+        trailing = _skip_horizontal_whitespace(source, end)
+        if trailing < len(source) and source[trailing] == ",":
+            trailing = _skip_horizontal_whitespace(source, trailing + 1)
+        if trailing >= len(source) or source[trailing] in "\r\n":
+            if source[trailing: trailing + 2] == "\r\n":
+                trailing += 2
+            elif trailing < len(source):
+                trailing += 1
+            return source[:line_start] + source[trailing:]
+
+    trailing = _skip_horizontal_whitespace(source, end)
+    if trailing < len(source) and source[trailing] == ",":
+        trailing = _skip_horizontal_whitespace(source, trailing + 1)
+        if trailing < len(source) and source[trailing] == " ":
+            trailing += 1
+        return source[:start] + source[trailing:]
+
+    leading = start
+    while leading > 0 and source[leading - 1] in " \t":
+        leading -= 1
+    if leading > 0 and source[leading - 1] == ",":
+        return source[: leading - 1] + source[end:]
+    return source[:start] + source[end:]
+
+
+def _skip_horizontal_whitespace(source: str, offset: int) -> int:
+    while offset < len(source) and source[offset] in " \t":
+        offset += 1
+    return offset
+
+
 def _read_builder_state(store: _ObjectStore, name: str) -> dict[str, Any]:
     try:
         data = json.loads(store.get(_builder_state_key(name)).decode("utf-8"))