Agent Admin

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Shared helpers for the agent-admin skill scripts.

Intentionally small: boto3 client factory, ISO8601 timestamps, UUID generation,
JSON output, fatal error exit, strict frontmatter `name` extractor, and a
db_schema write guard. All agent-admin scripts import from this module.
"""

from __future__ import annotations

import json
import re
import sys
import uuid
from datetime import datetime, timezone
from typing import Any

import boto3

# Exit codes (aligned with design.md)
EXIT_USAGE = 2
EXIT_NOT_FOUND = 3
EXIT_AWS_ERROR = 4
EXIT_IO_ERROR = 5
EXIT_FORBIDDEN = 6

# db_schema is intentionally excluded. No script writes it.
WRITABLE_FIELDS = frozenset({"system_prompt", "skills", "agent_name"})


def make_client(region: str) -> Any:
    """Return a low-level DynamoDB client in the given region."""
    return boto3.client("dynamodb", region_name=region)


def now_iso8601() -> str:
    """UTC timestamp like '2025-01-02T03:04:05+00:00' (matches existing data)."""
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def new_agent_id() -> str:
    """Lowercase uuid4 string."""
    return str(uuid.uuid4()).lower()


def print_json(obj: Any) -> None:
    """Print a single JSON object (UTF-8, indented) to stdout."""
    print(json.dumps(obj, ensure_ascii=False, indent=2))


def die(msg: str, code: int = 1) -> None:
    """Print ``msg`` to stderr and exit with ``code``."""
    print(msg, file=sys.stderr)
    sys.exit(code)


def warn(msg: str) -> None:
    """Print ``msg`` to stderr without exiting."""
    print(msg, file=sys.stderr)


def read_text_file(path: str) -> str:
    """Read a UTF-8 file; die with EXIT_IO_ERROR on failure."""
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except OSError as e:
        die(f"cannot read file: {path}: {e}", EXIT_IO_ERROR)
        raise  # unreachable; for type-checkers


# ---------------------------------------------------------------------------
# Frontmatter parser
# ---------------------------------------------------------------------------

_KEY_LINE_RE = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_-]*)\s*:\s*(.*)$")


def extract_knowledge_name(entry: str) -> str:
    """Extract the ``name`` field from a Knowledge entry.

    Strict grammar: the entry must start with ``---\\n``, be followed by one or
    more ``key: value`` lines, then a terminating ``---\\n``, then the body.
    Only the ``name`` field is extracted; other keys are ignored.

    Raises ``ValueError`` with a human-readable message on any deviation.
    """
    if not entry.startswith("---\n"):
        raise ValueError("knowledge frontmatter missing: entry must start with '---\\n'")

    # Skip the opening fence line.
    rest = entry[len("---\n"):]
    lines = rest.split("\n")

    name: str | None = None
    terminated = False
    for i, line in enumerate(lines):
        if line == "---":
            terminated = True
            break
        m = _KEY_LINE_RE.match(line)
        if not m:
            # Allow blank lines inside the header defensively.
            if line.strip() == "":
                continue
            raise ValueError(
                f"knowledge frontmatter malformed at line {i + 2}: {line!r}"
            )
        key, value = m.group(1), m.group(2)
        if key == "name":
            name = value.strip()

    if not terminated:
        raise ValueError("knowledge frontmatter unterminated: missing closing '---'")
    if name is None:
        raise ValueError("knowledge name missing from frontmatter")
    if name == "":
        raise ValueError("knowledge name empty")
    return name


# ---------------------------------------------------------------------------
# Defensive: no writer script should ever accept db_schema as a field.
# ---------------------------------------------------------------------------

def assert_writable_field(name: str) -> None:
    """Abort with EXIT_FORBIDDEN if ``name`` is not in WRITABLE_FIELDS."""
    if name not in WRITABLE_FIELDS:
        die(
            f"field '{name}' is not writable from this skill "
            "(db_schema is managed by the Admin UI only)",
            code=EXIT_FORBIDDEN,
        )


# ---------------------------------------------------------------------------
# DynamoDB helpers
# ---------------------------------------------------------------------------

def unwrap_ddb(item: dict[str, Any]) -> dict[str, Any]:
    """Convert a DynamoDB low-level item to plain Python.

    Only ``S`` (string) and ``L`` (list of S) are handled because that is all
    this table uses. Unknown types are passed through unchanged so unexpected
    attributes remain visible for debugging.
    """
    out: dict[str, Any] = {}
    for k, v in item.items():
        if not isinstance(v, dict) or len(v) != 1:
            out[k] = v
            continue
        (tag, val), = v.items()
        if tag == "S":
            out[k] = val
        elif tag == "L":
            out[k] = [elem.get("S", elem) for elem in val]
        else:
            out[k] = v
    return out

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Create a new agent in the agent-admin config table.

Usage:
    uv run create_agent.py \\
        --table-name TABLE --region REGION \\
        --agent-name NAME \\
        [--system-prompt-file PATH]

Inserts a new item with a freshly generated lowercase UUID.
``db_schema`` is deliberately left empty — the Agent Runtime falls back to the
``default`` agent's db_schema when a manual agent has none.
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    die,
    make_client,
    new_agent_id,
    now_iso8601,
    print_json,
    read_text_file,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Create a new agent.")
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--agent-name", required=True)
    parser.add_argument(
        "--system-prompt-file",
        default=None,
        help="Optional path to a UTF-8 file containing the system_prompt.",
    )
    args = parser.parse_args()

    agent_name = args.agent_name.strip()
    if not agent_name:
        from _common import EXIT_USAGE

        die("agent-name must not be empty", EXIT_USAGE)

    system_prompt = (
        read_text_file(args.system_prompt_file) if args.system_prompt_file else ""
    )

    agent_id = new_agent_id()
    now = now_iso8601()

    client = make_client(args.region)
    try:
        client.put_item(
            TableName=args.table_name,
            Item={
                "id": {"S": agent_id},
                "agent_name": {"S": agent_name},
                "system_prompt": {"S": system_prompt},
                "db_schema": {"S": ""},
                "skills": {"L": []},
                "created_at": {"S": now},
                "updated_at": {"S": now},
            },
            ConditionExpression="attribute_not_exists(id)",
        )
    except ClientError as e:
        err = e.response.get("Error", {})
        die(f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}", EXIT_AWS_ERROR)

    print_json({"id": agent_id, "agent_name": agent_name, "created_at": now})


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Remove one Knowledge entry by name.

Usage:
    uv run delete_knowledge.py \\
        --table-name TABLE --region REGION \\
        [--id ID] \\
        --name NAME

Walks the DynamoDB ``skills`` attribute, drops the entry whose frontmatter
``name`` matches ``--name``, and rewrites the attribute. If no entry matched,
exits with EXIT_NOT_FOUND.
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    EXIT_NOT_FOUND,
    die,
    extract_knowledge_name,
    make_client,
    now_iso8601,
    print_json,
    read_text_file,  # noqa: F401 (kept available for symmetry)
    unwrap_ddb,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Delete one Knowledge entry.")
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--id", default="default")
    parser.add_argument("--name", required=True)
    args = parser.parse_args()

    client = make_client(args.region)
    try:
        resp = client.get_item(
            TableName=args.table_name, Key={"id": {"S": args.id}}
        )
    except ClientError as e:
        err = e.response.get("Error", {})
        die(f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}", EXIT_AWS_ERROR)

    if "Item" not in resp:
        die(f"agent not found: {args.id}", EXIT_NOT_FOUND)

    existing = unwrap_ddb(resp["Item"]).get("skills", []) or []

    new_list: list[str] = []
    matched = False
    for entry in existing:
        if not isinstance(entry, str):
            new_list.append(entry)  # type: ignore[arg-type]
            continue
        try:
            entry_name = extract_knowledge_name(entry)
        except ValueError:
            entry_name = None
        if entry_name == args.name:
            matched = True
        else:
            new_list.append(entry)

    if not matched:
        die(f"knowledge not found: {args.name}", EXIT_NOT_FOUND)

    now = now_iso8601()
    try:
        client.update_item(
            TableName=args.table_name,
            Key={"id": {"S": args.id}},
            UpdateExpression="SET skills = :sk, updated_at = :ua",
            ExpressionAttributeValues={
                ":sk": {"L": [{"S": s} for s in new_list]},
                ":ua": {"S": now},
            },
            ConditionExpression="attribute_exists(id)",
        )
    except ClientError as e:
        code = e.response.get("Error", {}).get("Code", "")
        if code == "ConditionalCheckFailedException":
            die(f"agent not found: {args.id}", EXIT_NOT_FOUND)
        msg = e.response.get("Error", {}).get("Message", str(e))
        die(f"aws error: {code}: {msg}", EXIT_AWS_ERROR)

    print_json(
        {
            "id": args.id,
            "name": args.name,
            "knowledge_count": len(new_list),
            "updated_at": now,
        }
    )


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Fetch the full state of a single agent.

Usage:
    uv run get_agent.py --table-name TABLE --region REGION [--id ID]

Returns one JSON object with all relevant attributes, including ``db_schema``
and the full ``knowledge`` list. The DynamoDB attribute ``skills`` is renamed
to ``knowledge`` at the script boundary (Knowledge is the user-facing term).

db_schema fallback
------------------
If the target agent has an empty ``db_schema`` and is not itself the
``default`` agent, the script performs a second ``get_item`` against
``id="default"`` and returns that db_schema instead. The output always
carries ``db_schema_source: "self" | "default"`` so the caller knows where
the value came from. This mirrors the Agent Runtime behavior in
``agent/agent.py`` when loading config.
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    EXIT_NOT_FOUND,
    die,
    make_client,
    print_json,
    unwrap_ddb,
    warn,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Fetch one agent's full state.")
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--id", default="default")
    args = parser.parse_args()

    client = make_client(args.region)

    try:
        resp = client.get_item(
            TableName=args.table_name, Key={"id": {"S": args.id}}
        )
    except ClientError as e:
        err = e.response.get("Error", {})
        die(f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}", EXIT_AWS_ERROR)

    if "Item" not in resp:
        die(f"agent not found: {args.id}", EXIT_NOT_FOUND)

    item = unwrap_ddb(resp["Item"])

    db_schema = item.get("db_schema", "") or ""
    db_schema_source = "self"

    if db_schema == "" and args.id != "default":
        try:
            default_resp = client.get_item(
                TableName=args.table_name, Key={"id": {"S": "default"}}
            )
        except ClientError as e:
            err = e.response.get("Error", {})
            die(
                f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}",
                EXIT_AWS_ERROR,
            )
        if "Item" in default_resp:
            default_item = unwrap_ddb(default_resp["Item"])
            db_schema = default_item.get("db_schema", "") or ""
            db_schema_source = "default"
        else:
            warn("default agent not found; returning empty db_schema")
            db_schema = ""
            db_schema_source = "default"

    knowledge = item.get("skills", []) or []

    print_json(
        {
            "id": args.id,
            "agent_name": item.get("agent_name"),
            "system_prompt": item.get("system_prompt", ""),
            "db_schema": db_schema,
            "db_schema_source": db_schema_source,
            "knowledge": knowledge,
            "created_at": item.get("created_at"),
            "updated_at": item.get("updated_at"),
        }
    )


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""List all agents in the agent-admin config table.

Usage:
    uv run list_agents.py --table-name TABLE --region REGION

Outputs JSON of the shape:
    {"agents": [{"id": str, "agent_name": str|null, "created_at": str|null}, ...]}
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

# Allow running via `uv run <path>` from anywhere: add script dir to sys.path
# so the `_common` sibling module is importable.
sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    die,
    make_client,
    print_json,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="List agents in the config table.")
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    args = parser.parse_args()

    client = make_client(args.region)
    agents: list[dict] = []
    last_key: dict | None = None

    try:
        while True:
            kwargs = {
                "TableName": args.table_name,
                "ProjectionExpression": "id, agent_name, created_at",
            }
            if last_key is not None:
                kwargs["ExclusiveStartKey"] = last_key
            resp = client.scan(**kwargs)
            for item in resp.get("Items", []):
                agents.append(
                    {
                        "id": item["id"]["S"],
                        "agent_name": item.get("agent_name", {}).get("S"),
                        "created_at": item.get("created_at", {}).get("S"),
                    }
                )
            last_key = resp.get("LastEvaluatedKey")
            if last_key is None:
                break
    except ClientError as e:
        err = e.response.get("Error", {})
        die(f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}", EXIT_AWS_ERROR)

    print_json({"agents": agents})


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Replace an agent's ``system_prompt`` with the contents of a file.

Usage:
    uv run update_system_prompt.py \\
        --table-name TABLE --region REGION \\
        [--id ID] \\
        --prompt-file PATH

The script only accepts file input to sidestep shell-quoting hazards; Kiro is
expected to produce the file via its native ``fsWrite`` tool.
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    EXIT_NOT_FOUND,
    die,
    make_client,
    now_iso8601,
    print_json,
    read_text_file,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Update an agent's system_prompt.")
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--id", default="default")
    parser.add_argument("--prompt-file", required=True)
    args = parser.parse_args()

    prompt = read_text_file(args.prompt_file)
    now = now_iso8601()

    client = make_client(args.region)
    try:
        client.update_item(
            TableName=args.table_name,
            Key={"id": {"S": args.id}},
            UpdateExpression="SET system_prompt = :sp, updated_at = :ua",
            ExpressionAttributeValues={
                ":sp": {"S": prompt},
                ":ua": {"S": now},
            },
            ConditionExpression="attribute_exists(id)",
        )
    except ClientError as e:
        code = e.response.get("Error", {}).get("Code", "")
        if code == "ConditionalCheckFailedException":
            die(f"agent not found: {args.id}", EXIT_NOT_FOUND)
        msg = e.response.get("Error", {}).get("Message", str(e))
        die(f"aws error: {code}: {msg}", EXIT_AWS_ERROR)

    print_json(
        {
            "id": args.id,
            "updated_at": now,
            "bytes": len(prompt.encode("utf-8")),
        }
    )


if __name__ == "__main__":
    main()

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "boto3>=1.42",
# ]
# ///
"""Add or replace one Knowledge entry.

Usage:
    uv run upsert_knowledge.py \\
        --table-name TABLE --region REGION \\
        [--id ID] \\
        --knowledge-file PATH

The file must contain exactly one Knowledge entry in the form::

    ---
    name: <kebab-case>
    description: <short description>
    ---
    <markdown body>

If an existing entry with the same ``name`` is found in the DynamoDB
``skills`` attribute, it is replaced in place. Otherwise the new entry is
appended. The whole list is rewritten via ``update_item``, but the caller
only ever supplies one file — the merge happens inside this script.

Note on terminology: the DynamoDB attribute name remains ``skills`` for
backward compatibility with ``agent/agent.py`` and
``lambda/adminwebbackend/app.py``. At the script boundary everything is
called ``knowledge``.
"""

from __future__ import annotations

import argparse
import sys

from botocore.exceptions import ClientError

sys.path.insert(0, __file__.rsplit("/", 1)[0])
from _common import (  # noqa: E402
    EXIT_AWS_ERROR,
    EXIT_NOT_FOUND,
    EXIT_USAGE,
    die,
    extract_knowledge_name,
    make_client,
    now_iso8601,
    print_json,
    read_text_file,
    unwrap_ddb,
)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Add or replace one Knowledge entry."
    )
    parser.add_argument("--table-name", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--id", default="default")
    parser.add_argument("--knowledge-file", required=True)
    args = parser.parse_args()

    raw = read_text_file(args.knowledge_file)

    try:
        name = extract_knowledge_name(raw)
    except ValueError as e:
        die(str(e), EXIT_USAGE)

    client = make_client(args.region)
    try:
        resp = client.get_item(
            TableName=args.table_name, Key={"id": {"S": args.id}}
        )
    except ClientError as e:
        err = e.response.get("Error", {})
        die(f"aws error: {err.get('Code', '?')}: {err.get('Message', str(e))}", EXIT_AWS_ERROR)

    if "Item" not in resp:
        die(f"agent not found: {args.id}", EXIT_NOT_FOUND)

    existing = unwrap_ddb(resp["Item"]).get("skills", []) or []

    new_list: list[str] = []
    action = "added"
    matched = False
    for entry in existing:
        if not isinstance(entry, str):
            # Unexpected shape: pass through unchanged.
            new_list.append(entry)  # type: ignore[arg-type]
            continue
        try:
            entry_name = extract_knowledge_name(entry)
        except ValueError:
            entry_name = None
        if entry_name == name:
            new_list.append(raw)
            action = "updated"
            matched = True
        else:
            new_list.append(entry)
    if not matched:
        new_list.append(raw)

    now = now_iso8601()
    try:
        client.update_item(
            TableName=args.table_name,
            Key={"id": {"S": args.id}},
            UpdateExpression="SET skills = :sk, updated_at = :ua",
            ExpressionAttributeValues={
                ":sk": {"L": [{"S": s} for s in new_list]},
                ":ua": {"S": now},
            },
            ConditionExpression="attribute_exists(id)",
        )
    except ClientError as e:
        code = e.response.get("Error", {}).get("Code", "")
        if code == "ConditionalCheckFailedException":
            die(f"agent not found: {args.id}", EXIT_NOT_FOUND)
        msg = e.response.get("Error", {}).get("Message", str(e))
        die(f"aws error: {code}: {msg}", EXIT_AWS_ERROR)

    print_json(
        {
            "id": args.id,
            "name": name,
            "action": action,
            "knowledge_count": len(new_list),
            "updated_at": now,
        }
    )


if __name__ == "__main__":
    main()