Agentswarm/scripts/test-key-injection-contract.py

"""Model-key injection + revocation handshake contract tests (agent_swarm#16 / HM #60).

Pins the two Swarm-side parameters HM locked for #60 (calling-key injection) before HM writes its
mint→KV→revoke side:

  A.2 — KV secret VALUE format: Swarm resolves the per-user `sk-` from the HM #60 JSON
        ``{"openai_api_key": "sk-..."}`` (mirrors the callback-secret convention); a bare `sk-`
        string is still accepted; malformed / missing field never fabricates a key.
  A.5 — revocation signal: a `stop` is the only *final* state (completed/failed are reopenable via
        POST …/input, so they keep the key). When a stop drains a user's retained-run set, the
        runtime emits exactly one `swarm.pool_terminated{user_id, secret_ref}` so HM revokes the
        `sk-` + clears the KV secret. No double-revoke; still-live runs keep the key.

Hermetic: REDIS_FAKE, no model key, no real Manager callback (callback url empty).

Run from agent_swarm_v6 (install deps first — needs fakeredis):
    pip install -r orchestrator/requirements.txt
    REDIS_FAKE=1 python scripts/test-key-injection-contract.py
"""
import asyncio
import json
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator import agent_launcher
from orchestrator.swarm_runtime import swarm_runtime, FROZEN_CLIENT_EVENT_TYPES

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


async def stored_events(swarm_id):
    raw = await redis_client.lrange(f"{swarm_runtime.EVENT_KEY_PREFIX}{swarm_id}", 0, -1)
    return [json.loads(r) for r in raw]


async def new_run(user_id, secret_ref="azkv://heicode-vault/secrets/swarm-model-key-u1"):
    body = {
        "mode": "swarm",
        "orchestration_plan": {"objective": "key-contract test"},
        "billing_context": {"secret_ref": secret_ref},
        "callback": {"url": "", "subscribed_events": []},
        "metadata": {"manager_deployment_id": "m-key", "runtime_headers": {"x_user_id": user_id}},
    }
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="k")
    return run


# ---------------------------------------------------------------------------------------------
# A.2 — KV secret value format
# ---------------------------------------------------------------------------------------------
def test_kv_value_format():
    check("JSON {openai_api_key} → key extracted",
          agent_launcher._extract_model_key('{"openai_api_key": "sk-CONTRACT"}') == "sk-CONTRACT")
    check("JSON with sibling fields → still extracts the key",
          agent_launcher._extract_model_key('{"openai_api_key":"sk-X","note":"per-user"}') == "sk-X")
    check("bare sk- string value still accepted (back-compat)",
          agent_launcher._extract_model_key("sk-BARE") == "sk-BARE")
    check("malformed JSON → None (never fabricated)",
          agent_launcher._extract_model_key("{not json") is None)
    check("JSON missing the key field → None",
          agent_launcher._extract_model_key('{"callback_signing_secret":"x"}') is None)
    check("empty value → None", agent_launcher._extract_model_key("") is None)

    # End-to-end via resolve_model_key: azkv ref → dev env map → JSON value.
    os.environ.pop("AGENT_LAUNCH_MODEL_KEY", None)
    os.environ.pop("OPENAI_API_KEY", None)
    os.environ["HEICODE_SECRET_swarm-model-key-u1"] = '{"openai_api_key": "sk-RESOLVED"}'
    body = {"billing_context": {"secret_ref": "azkv://heicode-vault/secrets/swarm-model-key-u1"}}
    check("resolve_model_key resolves the azkv JSON secret to the bare sk-",
          agent_launcher.resolve_model_key(body) == "sk-RESOLVED")
    os.environ.pop("HEICODE_SECRET_swarm-model-key-u1", None)


# ---------------------------------------------------------------------------------------------
# A.5 — revocation handshake (swarm.pool_terminated on retained-set drain)
# ---------------------------------------------------------------------------------------------
async def pool_terminated_events(swarm_id):
    return [e for e in await stored_events(swarm_id) if e["event_type"] == "swarm.pool_terminated"]


async def test_single_run_stop_signals():
    run = await new_run("user-A")
    # not signalled before stop
    check("no pool_terminated before stop", not await pool_terminated_events(run.swarm_id))
    await swarm_runtime.stop_run(run.deployment_id, reason="manager stop")
    evs = await pool_terminated_events(run.swarm_id)
    check("pool_terminated emitted when user's only run is stopped", len(evs) == 1)
    if evs:
        p = evs[0]["payload"]
        check("pool_terminated carries user_id", p.get("user_id") == "user-A")
        check("pool_terminated carries azkv secret_ref (not a plaintext key)",
              str(p.get("secret_ref", "")).startswith("azkv://"))
    # It is an HM control-plane lifecycle signal, NOT a client task-cockpit event.
    check("pool_terminated is NOT in the frozen client event set",
          "swarm.pool_terminated" not in FROZEN_CLIENT_EVENT_TYPES)


async def test_multi_run_keeps_key_until_last_stop():
    a = await new_run("user-B")
    b = await new_run("user-B")
    await swarm_runtime.stop_run(a.deployment_id, reason="stop a")
    # B still retained → no signal on either run's stream yet
    sig_a = await pool_terminated_events(a.swarm_id)
    sig_b = await pool_terminated_events(b.swarm_id)
    check("stopping one of two runs does NOT revoke (other still live)", not sig_a and not sig_b)
    await swarm_runtime.stop_run(b.deployment_id, reason="stop b")
    check("pool_terminated fires only when the user's LAST run is stopped",
          len(await pool_terminated_events(b.swarm_id)) == 1)


async def test_no_double_revoke():
    run = await new_run("user-C")
    await swarm_runtime.stop_run(run.deployment_id, reason="stop 1")
    # a redundant second stop of the same (already drained) user must not signal again
    await swarm_runtime.stop_run(run.deployment_id, reason="stop 2")
    check("no double pool_terminated on a repeated stop",
          len(await pool_terminated_events(run.swarm_id)) == 1)


async def test_relife_after_new_run():
    # After a full teardown, a brand-new run re-arms the signal (clears the terminated flag).
    run1 = await new_run("user-D")
    await swarm_runtime.stop_run(run1.deployment_id, reason="stop d1")
    check("first teardown signals", len(await pool_terminated_events(run1.swarm_id)) == 1)
    run2 = await new_run("user-D")
    await swarm_runtime.stop_run(run2.deployment_id, reason="stop d2")
    check("a new run after teardown re-arms a second pool_terminated",
          len(await pool_terminated_events(run2.swarm_id)) == 1)


async def main():
    await redis_client.connect()
    test_kv_value_format()
    await test_single_run_stop_signals()
    await test_multi_run_keeps_key_until_last_stop()
    await test_no_double_revoke()
    await test_relife_after_new_run()
    print()
    if failures:
        print(f"{len(failures)} key-injection-contract check(s) FAILED: {failures}")
        sys.exit(1)
    print("all key-injection-contract checks passed")


if __name__ == "__main__":
    asyncio.run(main())