Merge branch 'main' into feat/agent-launcher-k8s-backend

scripts/test-key-injection-contract.py

@@ -0,0 +1,157 @@
+"""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-kv/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-kv/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())