{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-agent notepad — Python notebook\n",
    "\n",
    "Reproduces `demo.sh` end-to-end against a running OpenShell gateway, without shelling out to the `openshell` CLI. Every step uses the Python SDK at `python/openshell/`.\n",
    "\n",
    "**Prerequisites**\n",
    "\n",
    "- A running OpenShell gateway (`openshell gateway start` or already deployed).\n",
    "- A local Codex sign-in (`codex login`); the notebook reads tokens from `~/.codex/auth.json`.\n",
    "- `DEMO_GITHUB_OWNER`, `DEMO_GITHUB_REPO`, `DEMO_GITHUB_TOKEN` exported in the environment.\n",
    "- The `openshell` Python package installed (from the repo root: `uv pip install .`).\n",
    "\n",
    "Optional knobs (env vars, with defaults): `DEMO_TOPIC`, `DEMO_AGENT_COUNT` (1–8), `DEMO_BRANCH`, `DEMO_RUN_ID`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "import time\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from pathlib import Path\n",
    "\n",
    "import grpc\n",
    "\n",
    "from openshell import SandboxClient, SandboxError, policy_from_yaml\n",
    "\n",
    "# Resolve the demo directory whether the notebook runs from its own folder or repo root.\n",
    "_candidates = [Path.cwd(), Path.cwd() / \"examples\" / \"multi-agent-notepad\"]\n",
    "SCRIPT_DIR = next((c for c in _candidates if (c / \"runner.sh\").exists()), None)\n",
    "if SCRIPT_DIR is None:\n",
    "    raise SystemExit(\"Could not locate runner.sh; run from examples/multi-agent-notepad/ or the repo root.\")\n",
    "\n",
    "POLICY_TEMPLATE = SCRIPT_DIR / \"policy.template.yaml\"\n",
    "PROMPTS_DIR = SCRIPT_DIR / \"prompts\"\n",
    "RUNNER_SH = SCRIPT_DIR / \"runner.sh\"\n",
    "\n",
    "DEMO_TOPIC = os.environ.get(\"DEMO_TOPIC\", \"How should teams evaluate sandboxed coding agents?\")\n",
    "DEMO_AGENT_COUNT = int(os.environ.get(\"DEMO_AGENT_COUNT\", \"3\"))\n",
    "DEMO_BRANCH = os.environ.get(\"DEMO_BRANCH\", \"main\")\n",
    "DEMO_RUN_ID = os.environ.get(\"DEMO_RUN_ID\", time.strftime(\"%Y%m%d-%H%M%S\"))\n",
    "\n",
    "if not 1 <= DEMO_AGENT_COUNT <= 8:\n",
    "    raise SystemExit(\"DEMO_AGENT_COUNT must be between 1 and 8\")\n",
    "\n",
    "CODEX_PROVIDER = f\"codex-oauth-{DEMO_RUN_ID}\"\n",
    "GITHUB_PROVIDER = f\"github-memory-{DEMO_RUN_ID}\"\n",
    "\n",
    "WORKER_SLICES = [\n",
    "    \"Adoption criteria: what makes a sandboxed coding-agent workflow trustworthy enough to try?\",\n",
    "    \"Operational risks: what can go wrong when many autonomous agents run at once?\",\n",
    "    \"Security controls: which controls make the biggest difference for credential and network safety?\",\n",
    "    \"Developer experience: where does the workflow need to feel simple for users to adopt it?\",\n",
    "    \"Measurement: what signals show whether the agents produced useful work?\",\n",
    "    \"Scaling: what changes when the team runs dozens of agents instead of five?\",\n",
    "    \"Governance: what review and approval steps should remain human-controlled?\",\n",
    "    \"Repository hygiene: what makes the shared markdown notepad easy to review and clean up?\",\n",
    "]\n",
    "\n",
    "print(f\"Run id: {DEMO_RUN_ID}\")\n",
    "print(f\"Topic: {DEMO_TOPIC}\")\n",
    "print(f\"Agent count: {DEMO_AGENT_COUNT}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Validate the environment and connect to the gateway\n",
    "\n",
    "Mirrors `validate_env` in `demo.sh`: load Codex tokens, check GitHub vars, and confirm the gateway is reachable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "github_owner = os.environ[\"DEMO_GITHUB_OWNER\"]\n",
    "github_repo = os.environ[\"DEMO_GITHUB_REPO\"]\n",
    "_ = os.environ[\"DEMO_GITHUB_TOKEN\"]  # checked here so we fail fast\n",
    "\n",
    "codex_auth_path = Path.home() / \".codex\" / \"auth.json\"\n",
    "if not codex_auth_path.exists():\n",
    "    raise SystemExit(\"missing local Codex sign-in; run: codex login\")\n",
    "\n",
    "tokens = json.loads(codex_auth_path.read_text())[\"tokens\"]\n",
    "for key in (\"access_token\", \"refresh_token\", \"account_id\"):\n",
    "    if not tokens.get(key):\n",
    "        raise SystemExit(f\"local Codex sign-in is missing {key}; run: codex login\")\n",
    "\n",
    "os.environ[\"CODEX_AUTH_ACCESS_TOKEN\"] = tokens[\"access_token\"]\n",
    "os.environ[\"CODEX_AUTH_REFRESH_TOKEN\"] = tokens[\"refresh_token\"]\n",
    "os.environ[\"CODEX_AUTH_ACCOUNT_ID\"] = tokens[\"account_id\"]\n",
    "\n",
    "client = SandboxClient.from_active_cluster()\n",
    "health = client.health()\n",
    "print(f\"Gateway reachable. Server version: {health.version!r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Render the policy template into a `SandboxPolicy`\n",
    "\n",
    "The CLI uses `sed` to substitute `__OWNER__`, `__REPO__`, and `__RUN_ID__`; we do the same in Python and feed the result through `policy_from_yaml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rendered = (\n",
    "    POLICY_TEMPLATE.read_text()\n",
    "    .replace(\"__OWNER__\", github_owner)\n",
    "    .replace(\"__REPO__\", github_repo)\n",
    "    .replace(\"__RUN_ID__\", DEMO_RUN_ID)\n",
    ")\n",
    "\n",
    "policy = policy_from_yaml(rendered)\n",
    "print(f\"Loaded policy v{policy.version} with rule sets: {sorted(policy.network_policies.keys())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create provider records\n",
    "\n",
    "Provider records hold credentials by env-var name. The OpenShell sandbox supervisor injects a placeholder for each credential and the proxy substitutes the real value at the network boundary, so the agent never sees raw tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def safe_provider_delete(name: str) -> None:\n",
    "    try:\n",
    "        client.providers.delete(name)\n",
    "    except grpc.RpcError as exc:\n",
    "        if exc.code() != grpc.StatusCode.NOT_FOUND:\n",
    "            raise\n",
    "\n",
    "for stale in (CODEX_PROVIDER, GITHUB_PROVIDER):\n",
    "    safe_provider_delete(stale)\n",
    "\n",
    "client.providers.create(\n",
    "    name=CODEX_PROVIDER,\n",
    "    provider_type=\"generic\",\n",
    "    credentials_from_env=[\n",
    "        \"CODEX_AUTH_ACCESS_TOKEN\",\n",
    "        \"CODEX_AUTH_REFRESH_TOKEN\",\n",
    "        \"CODEX_AUTH_ACCOUNT_ID\",\n",
    "    ],\n",
    ")\n",
    "client.providers.create(\n",
    "    name=GITHUB_PROVIDER,\n",
    "    provider_type=\"generic\",\n",
    "    credentials_from_env=[\"DEMO_GITHUB_TOKEN\"],\n",
    ")\n",
    "print(f\"Created providers: {CODEX_PROVIDER}, {GITHUB_PROVIDER}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Stage the payload directory\n",
    "\n",
    "We mirror `demo.sh`'s layout: `payload/demo-runner.sh` and `payload/prompts/*.md`. The whole `payload/` directory is uploaded into `/sandbox` inside each worker, so the entrypoint runs `bash /sandbox/payload/demo-runner.sh ...`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir = Path(tempfile.mkdtemp(prefix=\"openshell-notebook-\"))\n",
    "payload_dir = work_dir / \"payload\"\n",
    "(payload_dir / \"prompts\").mkdir(parents=True)\n",
    "\n",
    "shutil.copy(RUNNER_SH, payload_dir / \"demo-runner.sh\")\n",
    "(payload_dir / \"demo-runner.sh\").chmod(0o755)\n",
    "for prompt in PROMPTS_DIR.glob(\"*.md\"):\n",
    "    shutil.copy(prompt, payload_dir / \"prompts\" / prompt.name)\n",
    "\n",
    "print(f\"Payload staged at {payload_dir}\")\n",
    "print(\"  files:\", sorted(p.relative_to(payload_dir).as_posix() for p in payload_dir.rglob(\"*\") if p.is_file()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Launch worker sandboxes in parallel\n",
    "\n",
    "Each worker gets its own isolated sandbox. The bash demo uses `&` + `wait`; we use `ThreadPoolExecutor` for the same fan-out — gRPC sync calls release the GIL, so each thread waits independently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_worker(index: int) -> str:\n",
    "    name = f\"codex-gh-{DEMO_RUN_ID}-a{index}\"\n",
    "    slice_text = WORKER_SLICES[(index - 1) % len(WORKER_SLICES)]\n",
    "    sandbox = client.create(\n",
    "        name=name,\n",
    "        image=\"base\",\n",
    "        providers=[CODEX_PROVIDER, GITHUB_PROVIDER],\n",
    "        policy=policy,\n",
    "    )\n",
    "    client.wait_ready(name)\n",
    "    client.upload(sandbox.id, payload_dir, \"/sandbox\")\n",
    "    result = client.exec(\n",
    "        sandbox.id,\n",
    "        [\n",
    "            \"bash\",\n",
    "            \"/sandbox/payload/demo-runner.sh\",\n",
    "            \"worker\",\n",
    "            github_owner,\n",
    "            github_repo,\n",
    "            DEMO_BRANCH,\n",
    "            DEMO_RUN_ID,\n",
    "            str(index),\n",
    "            str(DEMO_AGENT_COUNT),\n",
    "            DEMO_TOPIC,\n",
    "            slice_text,\n",
    "        ],\n",
    "    )\n",
    "    if result.exit_code != 0:\n",
    "        raise RuntimeError(\n",
    "            f\"worker {index} failed (exit {result.exit_code}); stderr tail:\\n{result.stderr[-1000:]}\"\n",
    "        )\n",
    "    return name\n",
    "\n",
    "\n",
    "worker_names: list[str] = []\n",
    "with ThreadPoolExecutor(max_workers=DEMO_AGENT_COUNT) as pool:\n",
    "    futures = {pool.submit(run_worker, i): i for i in range(1, DEMO_AGENT_COUNT + 1)}\n",
    "    for fut in as_completed(futures):\n",
    "        idx = futures[fut]\n",
    "        try:\n",
    "            worker_names.append(fut.result())\n",
    "            print(f\"  worker {idx} complete\")\n",
    "        except Exception as exc:\n",
    "            print(f\"  worker {idx} FAILED: {exc}\")\n",
    "\n",
    "if len(worker_names) != DEMO_AGENT_COUNT:\n",
    "    raise SystemExit(f\"only {len(worker_names)}/{DEMO_AGENT_COUNT} workers completed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run the synthesis sandbox\n",
    "\n",
    "Reads every worker's note from GitHub and writes the consolidated summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synth_name = f\"codex-gh-{DEMO_RUN_ID}-summary\"\n",
    "synth = client.create(\n",
    "    name=synth_name,\n",
    "    image=\"base\",\n",
    "    providers=[CODEX_PROVIDER, GITHUB_PROVIDER],\n",
    "    policy=policy,\n",
    ")\n",
    "client.wait_ready(synth_name)\n",
    "client.upload(synth.id, payload_dir, \"/sandbox\")\n",
    "result = client.exec(\n",
    "    synth.id,\n",
    "    [\n",
    "        \"bash\",\n",
    "        \"/sandbox/payload/demo-runner.sh\",\n",
    "        \"synthesis\",\n",
    "        github_owner,\n",
    "        github_repo,\n",
    "        DEMO_BRANCH,\n",
    "        DEMO_RUN_ID,\n",
    "        \"0\",\n",
    "        str(DEMO_AGENT_COUNT),\n",
    "        DEMO_TOPIC,\n",
    "    ],\n",
    ")\n",
    "if result.exit_code != 0:\n",
    "    raise SandboxError(f\"synthesis failed (exit {result.exit_code}); stderr:\\n{result.stderr[-1500:]}\")\n",
    "print(\"synthesis complete\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Print URLs and clean up\n",
    "\n",
    "Deletes worker + synthesis sandboxes and the demo's provider records, mirroring the `cleanup` trap in `demo.sh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def safe_sandbox_delete(name: str) -> None:\n",
    "    try:\n",
    "        client.delete(name)\n",
    "    except grpc.RpcError as exc:\n",
    "        if exc.code() != grpc.StatusCode.NOT_FOUND:\n",
    "            print(f\"  delete {name} failed: {exc.code()}\")\n",
    "\n",
    "for i in range(1, DEMO_AGENT_COUNT + 1):\n",
    "    safe_sandbox_delete(f\"codex-gh-{DEMO_RUN_ID}-a{i}\")\n",
    "safe_sandbox_delete(synth_name)\n",
    "\n",
    "for prov in (CODEX_PROVIDER, GITHUB_PROVIDER):\n",
    "    safe_provider_delete(prov)\n",
    "\n",
    "shutil.rmtree(work_dir, ignore_errors=True)\n",
    "client.close()\n",
    "\n",
    "base_url = f\"https://github.com/{github_owner}/{github_repo}/tree/{DEMO_BRANCH}/runs/{DEMO_RUN_ID}\"\n",
    "print(f\"\\nDone. Shared agent notepad: {base_url}\")\n",
    "print(f\"  summary:  {base_url}/summary.md\")\n",
    "for i in range(1, DEMO_AGENT_COUNT + 1):\n",
    "    print(f\"  agent {i}:  {base_url}/notes/agent-{i}.md\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}