clanker-saddle/test/bundles/advisor.py

"""The advisor backend bundle (goal 27, task 1): the AI Chat loop's
storage + LLM turn machinery, driven end-to-end against an in-process
scripted OpenAI-compatible chat-completions endpoint (the FakeSaddle
pattern). The scenarios:

- the migration: advisor_chats / advisor_messages / advisor_suggestions
  exist with the goal-27 shapes
- a full turn: user message → scripted tool call (list_projects, whose
  result rides the NEXT request's messages as a `tool` role row) →
  ask_question parks the chat ('awaiting_answer', an assistant row
  with the options payload, no further LLM call)
- answering: the option is recorded, "Human selected: X" lands as a
  user row, and the turn resumes to a plain assistant reply ('idle')
- suggestions: suggest_task inserts a PENDING row owned by the chat's
  user; nothing exists in tasks/goals yet
- the error path: one 500 from the endpoint appends a kind='error'
  message and the chat returns to idle
- the iteration cap: an endpoint that always answers with a tool call
  stops at 12 rounds with an error row, never looping forever
- privacy: a second user cannot see or turn the first user's chat
"""
from harness import *  # noqa: F401,F403 — the shared wire helpers


class ScriptedLlm:
    """In-process OpenAI-compatible chat-completions endpoint: every
    POST /chat/completions is inspected (messages + tools) and the
    next scripted response is served. `script` is a list of dicts:
    {"tools": [{id, name, arguments}]} or {"content": "..."}; an
    exhausted script repeats its last entry. `fail_next` makes the
    next N requests answer 500."""

    def __init__(self, script):
        self.script = list(script)
        self.cursor = 0
        self.requests = []
        self.fail_next = 0
        self.lock = threading.Lock()
        self.port = None
        self.server = None

    def serve(self):
        service = self

        class Handler(BaseHTTPRequestHandler):
            protocol_version = "HTTP/1.1"

            def log_message(self, fmt, *args):
                pass

            def do_POST(self):
                length = int(self.headers.get("Content-Length") or 0)
                raw = self.rfile.read(length)
                try:
                    payload = json.loads(raw or b"{}")
                except ValueError:
                    payload = {}
                with service.lock:
                    service.requests.append(payload)
                    fail = service.fail_next > 0
                    if fail:
                        service.fail_next -= 1
                    index = min(service.cursor, len(service.script) - 1)
                    step = service.script[index]
                    service.cursor += 1
                if fail:
                    body = b'{"error": "scripted failure"}'
                    self.send_response(500)
                else:
                    message = {"role": "assistant"}
                    if step.get("content") is not None:
                        message["content"] = step["content"]
                    calls = step.get("tools") or []
                    if calls:
                        message["content"] = step.get("content") or ""
                        message["tool_calls"] = [
                            {"id": call["id"], "type": "function",
                             "function": {"name": call["name"],
                                          "arguments": call["arguments"]}}
                            for call in calls]
                        finish = "tool_calls"
                    else:
                        finish = "stop"
                    envelope = {"choices": [{"index": 0, "finish_reason": finish,
                                             "message": message}],
                                "model": "scripted"}
                    body = json.dumps(envelope).encode()
                    self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.port = self.server.server_address[1]
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        return self.server

    def kill(self):
        if self.server is not None:
            self.server.shutdown()
            self.server.server_close()
            self.server = None


def advisor_api(auth, method, path, payload=None):
    """A login-gated advisor API call with the statum slot header."""
    request = urllib.request.Request(
        WEB + path, data=json.dumps(payload or {}).encode(), method=method)
    request.add_header("Content-Type", "application/json")
    request.add_header("X-Statum-Slot",
                       "key=%s; as_at=%s" % (auth["key"], auth["as_at"]))
    try:
        with urllib.request.urlopen(request, timeout=15) as response:
            return response.status, json.loads(response.read().decode() or "{}")
    except urllib.error.HTTPError as error:
        body = error.read().decode()
        try:
            return error.code, json.loads(body)
        except ValueError:
            return error.code, {"raw": body}


def chat_rows(chat_id, what="*"):
    conn = db()
    try:
        return conn.execute(
            "select %s from advisor_messages where chat_id = ? order by id"
            % what, (chat_id,)).fetchall()
    finally:
        conn.close()


def wait_status(chat_id, status, timeout=30):
    def read():
        conn = db()
        try:
            row = conn.execute(
                "select status from advisor_chats where id = ?",
                (chat_id,)).fetchone()
            return row and row[0] == status and status or False
        finally:
            conn.close()
    return wait_for(read, timeout=timeout,
                    what="chat %d -> %s" % (chat_id, status))


def run(ctx):
    auth, fake, fake_id, fake_token = ctx.auth, ctx.fake, ctx.fake_id, ctx.fake_token

    def db_rows(sql, params=()):
        conn = db()
        try:
            return conn.execute(sql, params).fetchall()
        finally:
            conn.close()

    print("== advisor: the migration ==")
    tables = {r[0] for r in db_rows(
        "select name from sqlite_master where type = 'table'")}
    check({"advisor_chats", "advisor_messages", "advisor_suggestions"} <= tables,
          "M0006 creates the three advisor tables on a fresh DB")
    chat_columns = [r[1] for r in db_rows("pragma table_info(advisor_chats)")]
    check(chat_columns == ["id", "user_id", "title", "status", "created_at",
                           "updated_at"],
          "advisor_chats carries the goal-27 columns (%s)" % chat_columns)
    message_columns = [r[1] for r in db_rows("pragma table_info(advisor_messages)")]
    check(message_columns == ["id", "chat_id", "kind", "content", "payload",
                              "answer", "answered_at", "created_at"],
          "advisor_messages carries the goal-27 columns (%s)" % message_columns)
    suggestion_columns = [r[1] for r in db_rows(
        "pragma table_info(advisor_suggestions)")]
    check(suggestion_columns == ["id", "chat_id", "user_id", "kind", "payload",
                                 "status", "task_id", "goal_id", "created_at",
                                 "decided_at", "delivered_at"],
          "advisor_suggestions carries the goal-27 columns (%s)"
          % suggestion_columns)

    print("== advisor: a full turn — tool call, question park ==")
    llm = ScriptedLlm([
        {"tools": [{"id": "call-1", "name": "list_projects",
                    "arguments": "{}"}]},
        {"tools": [{"id": "call-2", "name": "ask_question",
                    "arguments": json.dumps({
                        "question": "Which project should this target?",
                        "options": ["demo", "other-project"]})}]},
        {"content": "Understood — targeting demo."},
    ])
    llm.serve()
    # The advisor endpoint must be pointed at the scripted LLM: the
    # app restart below re-reads ADVISOR_LLM_BASE_URL.
    os.environ["ADVISOR_LLM_BASE_URL"] = "http://127.0.0.1:%d" % llm.port
    os.environ["ADVISOR_LLM_MODEL"] = "scripted"
    ctx.app.terminate()
    try:
        ctx.app.wait(timeout=10)
    except subprocess.TimeoutExpired:
        ctx.app.kill()
    ctx.app = start_app()
    ctx.auth = auth = login("admin", "adminpass123")

    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "Fleet review"})
    check(status == 200 and body.get("ok") and int(body.get("chat_id", 0)) > 0,
          "POST /api/advisor/chat creates the chat")
    chat_id = int(body["chat_id"])

    status, _ = advisor_api(auth, "POST", "/api/advisor/turn",
                            {"chat_id": chat_id,
                             "message": "Which project should I target?"})
    check(status == 200, "POST /api/advisor/turn accepts the message")

    # The tool call round: the request the LLM saw carries the tools
    # array and, on the SECOND call, the tool result as a `tool` role
    # message (the wire format rebuilt from rows).
    check(wait_for(lambda: len(llm.requests) >= 1 or False, timeout=15,
                   what="first LLM call") is not None,
          "the turn calls the LLM")
    first_request = llm.requests[0]
    tool_names = [t["function"]["name"] for t in first_request.get("tools") or []]
    check("list_projects" in tool_names and "ask_question" in tool_names
          and "suggest_task" in tool_names and "suggest_goal" in tool_names,
          "the request declares the advisor tool set")
    roles = [m["role"] for m in first_request.get("messages") or []]
    check(roles[0] == "system" and "read-only" in
          (first_request["messages"][0]["content"]),
          "the system prompt leads the messages")
    check(any(m.get("role") == "user"
              and m.get("content") == "Which project should I target?"
              for m in first_request["messages"]),
          "the user message rides the request")

    check(wait_for(lambda: len(llm.requests) >= 2 or False, timeout=15,
                   what="second LLM call") is not None,
          "the tool result triggers the next LLM call")
    second_request = llm.requests[1]
    tool_rows = [m for m in second_request["messages"] if m.get("role") == "tool"]
    check(len(tool_rows) == 1 and tool_rows[0]["tool_call_id"] == "call-1"
          and json.loads(tool_rows[0]["content"]).get("projects"),
          "the tool result rides back as a `tool` role message with the row data")
    assistant_call = [m for m in second_request["messages"]
                      if m.get("role") == "assistant" and m.get("tool_calls")]
    check(len(assistant_call) == 1
          and assistant_call[0]["tool_calls"][0]["id"] == "call-1",
          "the assistant tool_calls row is reconstructed on the wire")

    # ask_question parks: status awaiting_answer, one assistant row
    # carrying the question + options payload, and NO further LLM call.
    check(wait_status(chat_id, "awaiting_answer"),
          "ask_question parks the chat in awaiting_answer")
    calls_after_park = len(llm.requests)
    questions = [r for r in chat_rows(chat_id, "kind, content, payload")
                 if r[0] == "assistant" and r[2]
                 and json.loads(r[2]).get("question")]
    check(len(questions) == 1 and questions[0][1]
          == "Which project should this target?",
          "the question is stored as an assistant row with its payload")
    options = json.loads(questions[0][2]).get("options")
    check(options == ["demo", "other-project"],
          "the question payload carries the options (%s)" % options)
    time.sleep(0.5)
    check(len(llm.requests) == calls_after_park,
          "the parked turn makes no further LLM calls")

    print("== advisor: answering resumes the turn ==")
    status, body = advisor_api(auth, "POST", "/api/advisor/answer",
                               {"chat_id": chat_id, "answer": "demo"})
    check(status == 200 and body.get("ok"), "the answer is accepted")
    check(wait_status(chat_id, "idle"), "the resumed turn ends idle")
    check(wait_for(lambda: len(llm.requests) >= 3 or False, timeout=15,
                   what="third LLM call") is not None,
          "the answer resumes the LLM loop")
    answer_rows = [r for r in chat_rows(chat_id, "kind, content")
                   if r[1] == "Human selected: demo"]
    check(len(answer_rows) == 1 and answer_rows[0][0] == "user",
          'the option answer lands as "Human selected: demo" (user row)')
    check(any(r[0] == "assistant" and r[1] == "Understood — targeting demo."
              for r in chat_rows(chat_id, "kind, content")),
          "the final assistant reply is stored")
    third_request = llm.requests[2]
    check(any(m.get("role") == "user"
              and m.get("content") == "Human selected: demo"
              for m in third_request["messages"]),
          "the recorded answer rides the next request verbatim")
    answered = db_rows("select answer, answered_at from advisor_messages"
                       " where chat_id = ? and answer is not null", (chat_id,))
    check(len(answered) == 1 and answered[0][0] == "demo"
          and answered[0][1] is not None,
          "the question row keeps the raw answer and its timestamp")

    print("== advisor: free-text answer wording ==")
    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "free text"})
    free_chat = int(body["chat_id"])
    llm2 = ScriptedLlm([
        {"tools": [{"id": "f1", "name": "ask_question",
                    "arguments": json.dumps({"question": "Name it?",
                                             "options": []})}]},
        {"content": "Noted."},
    ])
    llm2.serve()
    os.environ["ADVISOR_LLM_BASE_URL"] = "http://127.0.0.1:%d" % llm2.port
    ctx.app.terminate()
    try:
        ctx.app.wait(timeout=10)
    except subprocess.TimeoutExpired:
        ctx.app.kill()
    ctx.app = start_app()
    ctx.auth = auth = login("admin", "adminpass123")
    advisor_api(auth, "POST", "/api/advisor/turn",
                {"chat_id": free_chat, "message": "suggest a name"})
    check(wait_status(free_chat, "awaiting_answer"),
          "an empty-options question parks too (free-text question)")
    status, _ = advisor_api(auth, "POST", "/api/advisor/answer",
                            {"chat_id": free_chat, "answer": "Zephyr"})
    check(wait_status(free_chat, "idle"), "the free-text answer resumes")
    check(any(r[1] == "Human answered: Zephyr"
              for r in chat_rows(free_chat, "kind, content")),
          'a non-option answer lands as "Human answered: ..."')
    llm2.kill()

    print("== advisor: suggestions insert pending rows ==")
    llm3 = ScriptedLlm([
        {"tools": [{"id": "s1", "name": "list_projects",
                    "arguments": "{}"}]},
        {"tools": [{"id": "s2", "name": "suggest_task",
                    "arguments": json.dumps({
                        "title": "Advisor smoke test",
                        "mandate": "Verify the advisor suggestion arc.",
                        "mode": "test",
                        "project_id": 1,
                        "priority": 4})}]},
        {"tools": [{"id": "s3", "name": "suggest_goal",
                    "arguments": json.dumps({
                        "title": "Advisor goal",
                        "brief": "Prove goal suggestions work.",
                        "project_id": 1,
                        "branch_name": "goal/advisor-smoke"})}]},
        {"content": "Two suggestions staged for your approval."},
    ])
    llm3.serve()
    os.environ["ADVISOR_LLM_BASE_URL"] = "http://127.0.0.1:%d" % llm3.port
    ctx.app.terminate()
    try:
        ctx.app.wait(timeout=10)
    except subprocess.TimeoutExpired:
        ctx.app.kill()
    ctx.app = start_app()
    ctx.auth = auth = login("admin", "adminpass123")
    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "suggestions"})
    sug_chat = int(body["chat_id"])
    advisor_api(auth, "POST", "/api/advisor/turn",
                {"chat_id": sug_chat, "message": "stage two suggestions"})
    check(wait_status(sug_chat, "idle"), "the suggestion turn completes")
    # the final assistant reply lands AFTER the status flip in some
    # interleavings — wait for the transcript itself to be complete.
    check(wait_for(lambda: [r for r in chat_rows(sug_chat, "kind, content")
                            if r[0] == "assistant"
                            and r[1] == "Two suggestions staged for your approval."]
                            or False, timeout=15, what="final reply") is not None,
          "the suggestion turn's final reply is stored")
    suggestions = db_rows(
        "select kind, payload, status, task_id, goal_id, user_id"
        " from advisor_suggestions where chat_id = ? order by id", (sug_chat,))
    check(len(suggestions) == 2, "two suggestion rows landed")
    check(suggestions[0][0] == "task" and suggestions[0][2] == "pending"
          and suggestions[0][3] == 0,
          "the task suggestion is PENDING with no task created yet")
    task_payload = json.loads(suggestions[0][1])
    check(task_payload.get("title") == "Advisor smoke test"
          and task_payload.get("mandate") == "Verify the advisor suggestion arc."
          and task_payload.get("mode") == "test"
          and task_payload.get("priority") == 4,
          "the task payload keeps every field verbatim")
    check(suggestions[1][0] == "goal" and suggestions[1][2] == "pending"
          and suggestions[1][4] == 0,
          "the goal suggestion is PENDING with no goal created yet")
    user_id = db_rows("select user_id from advisor_chats where id = ?",
                      (sug_chat,))[0][0]
    check(suggestions[0][5] == user_id and suggestions[1][5] == user_id,
          "suggestions are owned by the chat's user")
    check(db_rows("select count(*) from tasks"
                  " where title = 'Advisor smoke test'")[0][0] == 0,
          "NO task row exists before approval")
    # The approve flow itself is task 3 of the goal — this bundle only
    # proves the rows are inert and correctly owned.
    llm3.kill()

    print("== advisor: the error path ==")
    llm4 = ScriptedLlm([{"content": "never"}])
    llm4.serve()
    llm4.fail_next = 1
    os.environ["ADVISOR_LLM_BASE_URL"] = "http://127.0.0.1:%d" % llm4.port
    ctx.app.terminate()
    try:
        ctx.app.wait(timeout=10)
    except subprocess.TimeoutExpired:
        ctx.app.kill()
    ctx.app = start_app()
    ctx.auth = auth = login("admin", "adminpass123")
    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "errors"})
    err_chat = int(body["chat_id"])
    advisor_api(auth, "POST", "/api/advisor/turn",
                {"chat_id": err_chat, "message": "break once"})
    check(wait_status(err_chat, "idle"), "the errored turn returns the chat to idle")
    check(wait_for(lambda: [r for r in chat_rows(err_chat, "kind, content")
                            if r[0] == "error"] or False,
                   timeout=15, what="error row") is not None,
          "the failure appends its error message")
    error_rows = [r for r in chat_rows(err_chat, "kind, content")
                  if r[0] == "error"]
    check(len(error_rows) == 1 and "failed" in error_rows[0][1].lower(),
          "the failure appends an error message row (%s)"
          % [r[1] for r in error_rows])
    llm4.kill()

    print("== advisor: the iteration cap ==")
    llm5 = ScriptedLlm([
        {"tools": [{"id": "loop-%d" % index, "name": "list_projects",
                    "arguments": "{}"}]} for index in range(50)
    ])
    llm5.serve()
    os.environ["ADVISOR_LLM_BASE_URL"] = "http://127.0.0.1:%d" % llm5.port
    ctx.app.terminate()
    try:
        ctx.app.wait(timeout=10)
    except subprocess.TimeoutExpired:
        ctx.app.kill()
    ctx.app = start_app()
    ctx.auth = auth = login("admin", "adminpass123")
    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "loop"})
    loop_chat = int(body["chat_id"])
    advisor_api(auth, "POST", "/api/advisor/turn",
                {"chat_id": loop_chat, "message": "loop forever"})
    # The terminal state is the ERROR ROW (the status flips to idle
    # just before it lands) — waiting for that makes both the round
    # count and the row deterministic without sleeps.
    check(wait_for(lambda: [r for r in chat_rows(loop_chat, "kind")
                            if r[0] == "error"] or False,
                   timeout=30, what="capped turn error row") is not None,
          "the capped turn ends with an error row")
    check(wait_status(loop_chat, "idle"),
          "the capped turn terminates (chat back to idle)")
    check(len(llm5.requests) == 12,
          "the cap stopped further calls (exactly 12 rounds, %d seen)"
          % len(llm5.requests))
    llm5.kill()

    print("== advisor: per-user privacy ==")
    # A second user sees nothing of the first user's chats.
    admin_id = db_rows("select user_id from advisor_chats where id = ?",
                       (chat_id,))[0][0]
    second = "bob-" + str(int(time.time()))
    users_state = slot_public(auth, "/users", "users")
    act(auth, users_state["create"], {"username": second,
                                      "email": second + "@example.com",
                                      "forename": "Bob", "surname": "Advisor",
                                      "date_of_birth": "1990-05-05",
                                      "password": "bobpass123",
                                      "confirm_password": "bobpass123"})
    bob = login(second, "bobpass123")
    bob_id = db_rows("select id from spry_users where username = ?", (second,))[0][0]
    check(bob_id != admin_id, "the second user is distinct")
    status, body = advisor_api(bob, "GET", "/api/advisor/chat/%d" % chat_id)
    check(status == 404, "another user cannot read the chat (404)")
    status, body = advisor_api(bob, "POST", "/api/advisor/turn",
                               {"chat_id": chat_id, "message": "hi"})
    check(status == 404, "another user cannot turn the chat")
    status, body = advisor_api(bob, "POST", "/api/advisor/chat/delete",
                               {"chat_id": chat_id})
    check(status == 200 and db_rows(
        "select count(*) from advisor_chats where id = ?", (chat_id,))[0][0] == 1,
        "another user's delete is a no-op for the row")

    print("== advisor: chat list + delete ==")
    status, body = advisor_api(auth, "POST", "/api/advisor/chat",
                               {"title": "to delete"})
    doomed = int(body["chat_id"])
    advisor_api(auth, "POST", "/api/advisor/turn",
                {"chat_id": doomed, "message": "populate me"})
    # the turn still runs against llm5 (the loop server): wait for the
    # terminal error row, not merely the status flip, so delete cannot
    # race the final append.
    check(wait_for(lambda: [r for r in chat_rows(doomed, "kind")
                            if r[0] == "error"] or False,
                   timeout=30, what="populate turn to finish") is not None,
          "the populate turn finishes before the delete")
    status, _ = advisor_api(auth, "POST", "/api/advisor/chat/delete",
                            {"chat_id": doomed})
    check(status == 200 and db_rows(
        "select count(*) from advisor_chats where id = ?", (doomed,))[0][0] == 0,
        "delete removes the chat")
    check(db_rows("select count(*) from advisor_messages where chat_id = ?",
                  (doomed,))[0][0] == 0,
        "delete cascades the messages")

    # Restore the daemon-facing LLM env the other checks may rely on
    # (the advisor env is test-local).
    os.environ.pop("ADVISOR_LLM_BASE_URL", None)
    os.environ.pop("ADVISOR_LLM_MODEL", None)