basant307/AI_Governance_Project

1#!/usr/bin/env python32"""Run real end-to-end smoke tests against qwen CLI using qwen_code_sdk.3 4This script is intentionally lightweight and avoids any test doubles.5It is useful for manual verification after changing SDK runtime behavior.6"""7 8from __future__ import annotations9 10import sys11 12if sys.version_info < (3, 10):  # noqa: UP03613    import json14 15    version = ".".join(str(part) for part in sys.version_info[:3])16    payload = {17        "ok": False,18        "stage": "startup",19        "error": f"Python >=3.10 is required, current version is {version}",20        "error_type": "RuntimeError",21    }22    print(json.dumps(payload, ensure_ascii=False, indent=2))23    raise SystemExit(2)24 25import argparse26import asyncio27import json28import subprocess29import threading30from collections.abc import AsyncIterator, Awaitable, Callable31from dataclasses import asdict, dataclass32from pathlib import Path33from queue import Empty, Queue34from typing import Any, TypeVar35 36SDK_ROOT = Path(__file__).resolve().parents[1]37SRC_ROOT = SDK_ROOT / "src"38if str(SRC_ROOT) not in sys.path:39    sys.path.insert(0, str(SRC_ROOT))40 41from qwen_code_sdk import (  # noqa: E40242    SDKUserMessage,43    SyncQuery,44    is_sdk_assistant_message,45    is_sdk_result_message,46    is_sdk_system_message,47    query,48    query_sync,49)50from qwen_code_sdk.transport import prepare_spawn_info  # noqa: E40251 52T = TypeVar("T")53 54 55@dataclass56class AsyncSingleResult:57    ok: bool58    assistant_text: str | None59    result_text: str | None60    session_id: str61 62 63@dataclass64class AsyncControlResult:65    ok: bool66    supported_commands_type: str67    saw_system_message: bool68    saw_result_message: bool69    session_id: str70 71 72@dataclass73class SyncResult:74    ok: bool75    saw_result_message: bool76    result_text: str | None77    session_id: str78 79 80def parse_args() -> argparse.Namespace:81    parser = argparse.ArgumentParser(82        description="Run real qwen_code_sdk smoke tests using qwen CLI",83    )84    parser.add_argument(85        "--qwen",86        default="qwen",87        help="Path or command for qwen executable (default: qwen)",88    )89    parser.add_argument(90        "--cwd",91        default=str(Path.cwd()),92        help="Working directory passed to SDK query options",93    )94    parser.add_argument(95        "--model",96        default=None,97        help="Optional model name. If set, script will call set_model(model).",98    )99    parser.add_argument(100        "--timeout-seconds",101        type=float,102        default=90.0,103        help="Timeout used for control/callback/stream-close options",104    )105    parser.add_argument(106        "--json-only",107        action="store_true",108        help="Print only JSON result (no progress logs)",109    )110    return parser.parse_args()111 112 113def check_qwen_cli_available(qwen_cmd: str, timeout_seconds: float) -> str:114    spawn_info = prepare_spawn_info(qwen_cmd)115    completed = subprocess.run(116        [spawn_info.command, *spawn_info.args, "--version"],117        check=True,118        capture_output=True,119        text=True,120        timeout=timeout_seconds,121    )122    return completed.stdout.strip()123 124 125def build_options(args: argparse.Namespace) -> dict[str, Any]:126    return {127        "cwd": args.cwd,128        "path_to_qwen_executable": args.qwen,129        "permission_mode": "yolo",130        "max_session_turns": 1,131        "timeout": {132            "control_request": args.timeout_seconds,133            "can_use_tool": args.timeout_seconds,134            "stream_close": args.timeout_seconds,135        },136    }137 138 139def extract_assistant_text(message: dict[str, Any]) -> str:140    content = message["message"].get("content", [])141    if not isinstance(content, list):142        return ""143 144    text_parts: list[str] = []145    for block in content:146        if isinstance(block, dict) and block.get("type") == "text":147            text_parts.append(str(block.get("text", "")))148    return "".join(text_parts)149 150 151async def run_async_single(args: argparse.Namespace) -> AsyncSingleResult:152    token = "SDK_REAL_ASYNC_OK"153    options = build_options(args)154    q = query(155        f"Reply exactly with {token}",156        options,157    )158 159    assistant_text: str | None = None160    result_text: str | None = None161    try:162        async for message in q:163            if is_sdk_assistant_message(message):164                assistant_text = (assistant_text or "") + extract_assistant_text(165                    message166                )167            if is_sdk_result_message(message):168                result_text = str(message.get("result", ""))169    finally:170        await q.close()171 172    ok = token in (assistant_text or "") and token in (result_text or "")173    return AsyncSingleResult(174        ok=ok,175        assistant_text=assistant_text,176        result_text=result_text,177        session_id=q.get_session_id(),178    )179 180 181async def run_async_controls(args: argparse.Namespace) -> AsyncControlResult:182    token = "SDK_REAL_CONTROL_OK"183    options = build_options(args)184    release_prompt = asyncio.Event()185 186    async def prompts() -> AsyncIterator[SDKUserMessage]:187        await release_prompt.wait()188        yield {189            "type": "user",190            "session_id": "00000000-0000-4000-8000-000000000001",191            "message": {192                "role": "user",193                "content": f"Reply exactly with {token}",194            },195            "parent_tool_use_id": None,196        }197 198    q = query(prompts(), options)199 200    supported: dict[str, Any] | None = None201    saw_system_message = False202    saw_result_message = False203    try:204        supported = await q.supported_commands()205        await q.set_permission_mode("plan")206        await q.set_permission_mode("yolo")207        if args.model:208            await q.set_model(args.model)209 210        release_prompt.set()211        async for message in q:212            if is_sdk_system_message(message):213                saw_system_message = True214            if is_sdk_result_message(message):215                saw_result_message = True216                break217    finally:218        await q.close()219 220    ok = isinstance(supported, dict) and saw_result_message221    return AsyncControlResult(222        ok=ok,223        supported_commands_type=type(supported).__name__,224        saw_system_message=saw_system_message,225        saw_result_message=saw_result_message,226        session_id=q.get_session_id(),227    )228 229 230async def run_stage(stage: str, coro: Awaitable[T], timeout_seconds: float) -> T:231    try:232        return await asyncio.wait_for(coro, timeout=timeout_seconds)233    except TimeoutError as exc:234        message = f"{stage} timed out after {timeout_seconds} seconds"235        raise TimeoutError(message) from exc236 237 238def run_sync(239    args: argparse.Namespace,240    on_query: Callable[[SyncQuery], None] | None = None,241) -> SyncResult:242    token = "SDK_REAL_SYNC_OK"243    options = build_options(args)244    q = query_sync(245        f"Reply exactly with {token}",246        options,247    )248    if on_query is not None:249        on_query(q)250 251    saw_result_message = False252    result_text: str | None = None253    try:254        for message in q:255            if is_sdk_result_message(message):256                saw_result_message = True257                result_text = str(message.get("result", ""))258                break259    finally:260        q.close()261 262    ok = saw_result_message and token in (result_text or "")263    return SyncResult(264        ok=ok,265        saw_result_message=saw_result_message,266        result_text=result_text,267        session_id=q.get_session_id(),268    )269 270 271def run_sync_with_timeout(args: argparse.Namespace) -> SyncResult:272    result_queue: Queue[SyncResult | BaseException] = Queue(maxsize=1)273    query_holder: dict[str, SyncQuery] = {}274 275    def remember_query(q: SyncQuery) -> None:276        query_holder["query"] = q277 278    def worker() -> None:279        try:280            result_queue.put(run_sync(args, on_query=remember_query))281        except BaseException as exc:282            result_queue.put(exc)283 284    thread = threading.Thread(285        target=worker,286        name="qwen-sdk-real-smoke-sync",287        daemon=True,288    )289    thread.start()290 291    try:292        item = result_queue.get(timeout=args.timeout_seconds)293    except Empty as exc:294        q = query_holder.get("query")295        if q is not None:296            q.close()297        raise TimeoutError(298            f"sync check timed out after {args.timeout_seconds} seconds"299        ) from exc300 301    thread.join(timeout=1.0)302    if isinstance(item, BaseException):303        raise item304    return item305 306 307def build_failure_payload(308    *,309    stage: str,310    exc: BaseException,311    qwen_version: str | None = None,312    completed: dict[str, Any] | None = None,313) -> dict[str, Any]:314    payload: dict[str, Any] = {315        "ok": False,316        "stage": stage,317        "error": str(exc),318        "error_type": type(exc).__name__,319    }320    if qwen_version is not None:321        payload["qwen_version"] = qwen_version322    if completed:323        payload["completed"] = completed324    return payload325 326 327async def main() -> int:328    args = parse_args()329 330    try:331        qwen_version = check_qwen_cli_available(args.qwen, args.timeout_seconds)332    except (subprocess.CalledProcessError, OSError, subprocess.TimeoutExpired) as exc:333        payload = build_failure_payload(stage="preflight", exc=exc)334        print(json.dumps(payload, ensure_ascii=False, indent=2))335        return 2336 337    stage = "async single-turn check"338    completed: dict[str, Any] = {}339    try:340        if not args.json_only:341            print(f"[smoke] qwen version: {qwen_version}")342            print(f"[smoke] running {stage}...")343        async_single = await run_stage(344            stage,345            run_async_single(args),346            args.timeout_seconds,347        )348        completed["async_single"] = asdict(async_single)349 350        stage = "async control check"351        if not args.json_only:352            print(f"[smoke] running {stage}...")353        async_controls = await run_stage(354            stage,355            run_async_controls(args),356            args.timeout_seconds,357        )358        completed["async_controls"] = asdict(async_controls)359 360        stage = "sync check"361        if not args.json_only:362            print(f"[smoke] running {stage}...")363        sync_result = run_sync_with_timeout(args)364        completed["sync"] = asdict(sync_result)365    except Exception as exc:366        payload = build_failure_payload(367            stage=stage,368            exc=exc,369            qwen_version=qwen_version,370            completed=completed,371        )372        print(json.dumps(payload, ensure_ascii=False, indent=2))373        return 1374 375    all_ok = async_single.ok and async_controls.ok and sync_result.ok376    payload = {377        "ok": all_ok,378        "qwen_version": qwen_version,379        "async_single": asdict(async_single),380        "async_controls": asdict(async_controls),381        "sync": asdict(sync_result),382    }383    print(json.dumps(payload, ensure_ascii=False, indent=2))384    return 0 if all_ok else 1385 386 387if __name__ == "__main__":388    raise SystemExit(asyncio.run(main()))389