add delete button

prompt/api.py

@@ -97,3 +97,45 @@ def classify_batch(request, task_id: Optional[int] = None, force: bool = False):
     threading.Thread(target=classify_messages_batch, args=(ids,), daemon=True).start()
 
     return {"message": f"开始分类 {len(ids)} 条消息", "count": len(ids)}
+
+
+@router.delete("/messages/{message_id}/pair")
+@login_required
+def delete_message_pair(request, message_id: int):
+    """
+    Delete a message pair (assistant message + preceding user message) and
+    any linked submission. Only the conversation owner can do this.
+    """
+    asst_msg = get_object_or_404(Message, id=message_id, role="assistant")
+
+    if asst_msg.conversation.user != request.user:
+        raise HttpError(403, "只能删除自己的消息")
+
+    # Find the preceding user message
+    user_msg = (
+        Message.objects.filter(
+            conversation=asst_msg.conversation,
+            created__lt=asst_msg.created,
+            role="user",
+        )
+        .order_by("-created")
+        .first()
+    )
+
+    # Delete messages first, then submission
+    submission_id = asst_msg.submission_id  # capture before deletion nulls it
+
+    if user_msg:
+        user_msg.delete()
+    asst_msg.delete()
+
+    submission_deleted = False
+    if submission_id:
+        from submission.models import Submission as SubmissionModel
+        try:
+            SubmissionModel.objects.filter(id=submission_id).delete()
+            submission_deleted = True
+        except Exception:
+            pass
+
+    return {"deleted": True, "submission_deleted": submission_deleted}

prompt/consumers.py

@@ -17,7 +17,6 @@ class PromptConsumer(AsyncWebsocketConsumer):
         self.current_user_message = None
         await self.accept()
 
-        # Load or create conversation, send history
         self.conversation = await self.get_or_create_conversation()
         history = await self.get_history()
         await self.send(text_data=json.dumps({
@@ -43,14 +42,11 @@ class PromptConsumer(AsyncWebsocketConsumer):
         if not prompt:
             return
 
-        # Save user message
         self.current_user_message = await self.save_message("user", prompt)
 
         try:
-            # Build history for LLM
             history = await self.get_history_for_llm()
 
-            # Stream AI response
             full_response = ""
             try:
                 async for chunk in stream_chat(history, model=model):
@@ -66,15 +62,14 @@ class PromptConsumer(AsyncWebsocketConsumer):
                 }))
                 return
 
-            # Extract code and save assistant message
             code = extract_code(full_response)
-            await self.save_message("assistant", full_response, code)
+            assistant_msg = await self.save_message("assistant", full_response, code)
             self.current_user_message = None
 
-            # Send completion with extracted code
             await self.send(text_data=json.dumps({
                 "type": "complete",
                 "code": code,
+                "message_id": assistant_msg.id,
             }))
 
         finally:
@@ -114,6 +109,7 @@ class PromptConsumer(AsyncWebsocketConsumer):
         messages = self.conversation.messages.filter(source="conversation")
         return [
             {
+                "id": m.id,
                 "role": m.role,
                 "content": m.content,
                 "code": {

prompt/migrations/0005_message_add_submission_fk.py

@@ -0,0 +1,20 @@
+# Generated by Django 6.0.1 on 2026-04-15 07:50
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('prompt', '0004_update_message_source_default'),
+        ('submission', '0010_remove_conversation_fk'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='message',
+            name='submission',
+            field=models.OneToOneField(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='source_message', to='submission.submission'),
+        ),
+    ]

prompt/models.py

@@ -32,6 +32,13 @@ class Message(models.Model):
     prompt_level = models.IntegerField(
         null=True, blank=True, default=None, db_index=True, verbose_name="提示词层级"
     )
+    submission = models.OneToOneField(
+        "submission.Submission",
+        null=True,
+        blank=True,
+        on_delete=models.SET_NULL,
+        related_name="source_message",
+    )
 
     class Meta:
         ordering = ("created",)

prompt/tests.py

@@ -0,0 +1,177 @@
+from django.test import TestCase
+from django.contrib.auth import get_user_model
+from task.models import Task
+from submission.models import Submission
+from prompt.models import Conversation, Message
+
+User = get_user_model()
+
+
+def _make_user(username="student1"):
+    return User.objects.create_user(username=username, password="pw")
+
+
+def _make_task():
+    return Task.objects.create(
+        title="Test Task", task_type="challenge", display=1, content=""
+    )
+
+
+class SubmissionMessageLinkTest(TestCase):
+    def setUp(self):
+        self.user = _make_user("student1")
+        self.task = _make_task()
+        self.conv = Conversation.objects.create(user=self.user, task=self.task)
+        self.user_msg = Message.objects.create(
+            conversation=self.conv, role="user", content="帮我做个按钮"
+        )
+        self.asst_msg = Message.objects.create(
+            conversation=self.conv, role="assistant", content="好的",
+            code_html="<button>OK</button>", code_css="", code_js=""
+        )
+
+    def test_create_submission_links_message(self):
+        """POST /api/submission/ with message_id links the assistant message"""
+        self.client.force_login(self.user)
+        resp = self.client.post(
+            "/api/submission/",
+            data={
+                "task_id": self.task.id,
+                "html": "<button>OK</button>",
+                "css": "",
+                "js": "",
+                "message_id": self.asst_msg.id,
+            },
+            content_type="application/json",
+        )
+        self.assertEqual(resp.status_code, 200)
+        self.asst_msg.refresh_from_db()
+        sub = Submission.objects.filter(user=self.user, task=self.task).first()
+        self.assertIsNotNone(sub)
+        self.assertEqual(self.asst_msg.submission, sub)
+
+    def test_create_submission_without_message_id(self):
+        """POST /api/submission/ without message_id still works"""
+        self.client.force_login(self.user)
+        resp = self.client.post(
+            "/api/submission/",
+            data={"task_id": self.task.id, "html": "<p>hi</p>", "css": "", "js": ""},
+            content_type="application/json",
+        )
+        self.assertEqual(resp.status_code, 200)
+
+    def test_cannot_link_message_from_different_task(self):
+        """message_id from a different task is silently ignored"""
+        other_task = Task.objects.create(
+            title="Other Task", task_type="challenge", display=2, content=""
+        )
+        other_conv = Conversation.objects.create(user=self.user, task=other_task)
+        other_msg = Message.objects.create(
+            conversation=other_conv, role="assistant", content="other"
+        )
+        self.client.force_login(self.user)
+        resp = self.client.post(
+            "/api/submission/",
+            data={
+                "task_id": self.task.id,
+                "html": "<p>x</p>",
+                "css": "",
+                "js": "",
+                "message_id": other_msg.id,
+            },
+            content_type="application/json",
+        )
+        self.assertEqual(resp.status_code, 200)
+        # Submission created, but message NOT linked (wrong task)
+        other_msg.refresh_from_db()
+        self.assertIsNone(other_msg.submission)
+
+
+class DeleteSubmissionCascadeTest(TestCase):
+    def setUp(self):
+        self.user = _make_user("student2")
+        self.task = _make_task()
+        self.conv = Conversation.objects.create(user=self.user, task=self.task)
+        self.user_msg = Message.objects.create(
+            conversation=self.conv, role="user", content="问题"
+        )
+        self.asst_msg = Message.objects.create(
+            conversation=self.conv, role="assistant", content="回答",
+            code_html="<p>hi</p>", code_css="", code_js=""
+        )
+        self.sub = Submission.objects.create(
+            user=self.user, task=self.task, html="<p>hi</p>", css="", js=""
+        )
+        self.asst_msg.submission = self.sub
+        self.asst_msg.save(update_fields=["submission"])
+
+    def test_delete_submission_also_deletes_message_pair(self):
+        """DELETE /api/submission/{id} deletes linked user+assistant messages"""
+        self.client.force_login(self.user)
+        resp = self.client.delete(f"/api/submission/{self.sub.id}")
+        self.assertEqual(resp.status_code, 200)
+        self.assertFalse(Submission.objects.filter(id=self.sub.id).exists())
+        self.assertFalse(Message.objects.filter(id=self.asst_msg.id).exists())
+        self.assertFalse(Message.objects.filter(id=self.user_msg.id).exists())
+
+    def test_delete_submission_without_linked_message(self):
+        """DELETE /api/submission/{id} works even with no linked message"""
+        sub2 = Submission.objects.create(
+            user=self.user, task=self.task, html="", css="", js=""
+        )
+        self.client.force_login(self.user)
+        resp = self.client.delete(f"/api/submission/{sub2.id}")
+        self.assertEqual(resp.status_code, 200)
+        self.assertFalse(Submission.objects.filter(id=sub2.id).exists())
+
+
+class DeleteMessagePairTest(TestCase):
+    def setUp(self):
+        self.user = _make_user("student3")
+        self.task = _make_task()
+        self.conv = Conversation.objects.create(user=self.user, task=self.task)
+        self.user_msg = Message.objects.create(
+            conversation=self.conv, role="user", content="问题"
+        )
+        self.asst_msg = Message.objects.create(
+            conversation=self.conv, role="assistant", content="回答",
+            code_html="<p>ok</p>", code_css="", code_js=""
+        )
+        self.sub = Submission.objects.create(
+            user=self.user, task=self.task, html="<p>ok</p>", css="", js=""
+        )
+        self.asst_msg.submission = self.sub
+        self.asst_msg.save(update_fields=["submission"])
+
+    def test_delete_message_pair_also_deletes_submission(self):
+        self.client.force_login(self.user)
+        resp = self.client.delete(f"/api/prompt/messages/{self.asst_msg.id}/pair")
+        self.assertEqual(resp.status_code, 200)
+        self.assertFalse(Message.objects.filter(id=self.asst_msg.id).exists())
+        self.assertFalse(Message.objects.filter(id=self.user_msg.id).exists())
+        self.assertFalse(Submission.objects.filter(id=self.sub.id).exists())
+        data = resp.json()
+        self.assertTrue(data["deleted"])
+        self.assertTrue(data["submission_deleted"])
+
+    def test_delete_message_pair_without_submission(self):
+        # Create a second pair without a submission link
+        user2 = Message.objects.create(
+            conversation=self.conv, role="user", content="另一问"
+        )
+        asst2 = Message.objects.create(
+            conversation=self.conv, role="assistant", content="另一条"
+        )
+        self.client.force_login(self.user)
+        resp = self.client.delete(f"/api/prompt/messages/{asst2.id}/pair")
+        self.assertEqual(resp.status_code, 200)
+        data = resp.json()
+        self.assertFalse(data["submission_deleted"])
+        self.assertFalse(Message.objects.filter(id=asst2.id).exists())
+        self.assertFalse(Message.objects.filter(id=user2.id).exists())
+
+    def test_delete_message_pair_forbidden_for_other_user(self):
+        other = _make_user("other")
+        self.client.force_login(other)
+        resp = self.client.delete(f"/api/prompt/messages/{self.asst_msg.id}/pair")
+        self.assertEqual(resp.status_code, 403)

submission/api.py

@@ -75,7 +75,7 @@ def create_submission(request, payload: SubmissionIn):
             from .classifier import classify_conversation_messages
             threading.Thread(target=classify_conversation_messages, args=(conversation.id,), daemon=True).start()
 
-    Submission.objects.create(
+    submission = Submission.objects.create(
         user=request.user,
         task=task,
         html=payload.html,
@@ -83,6 +83,22 @@ def create_submission(request, payload: SubmissionIn):
         js=payload.js,
     )
 
+    # Link assistant message if provided
+    if payload.message_id:
+        try:
+            msg = Message.objects.get(
+                id=payload.message_id,
+                role="assistant",
+                conversation__user=request.user,
+                conversation__task=task,
+            )
+            msg.submission = submission
+            msg.save(update_fields=["submission"])
+        except Message.DoesNotExist:
+            pass  # invalid message_id — submission already created, silently skip
+
+    return {"id": str(submission.id)}
+
 
 @router.get("/", response=List[SubmissionOut])
 @paginate
@@ -193,6 +209,23 @@ def delete_submission(request, submission_id: UUID):
     submission = get_object_or_404(Submission, id=submission_id)
     if submission.user != request.user:
         raise HttpError(403, "只能删除自己的提交")
+
+    # Delete linked message pair if present
+    asst_msg = Message.objects.filter(submission=submission).first()
+    if asst_msg:
+        user_msg = (
+            Message.objects.filter(
+                conversation=asst_msg.conversation,
+                created__lt=asst_msg.created,
+                role="user",
+            )
+            .order_by("-created")
+            .first()
+        )
+        if user_msg:
+            user_msg.delete()
+        asst_msg.delete()
+
     submission.delete()
     return {"message": "删除成功"}
 

submission/schemas.py

@@ -9,6 +9,7 @@ class SubmissionIn(Schema):
     css: Optional[str] = None
     js: Optional[str] = None
     prompt: Optional[str] = None
+    message_id: Optional[int] = None  # 关联的 assistant message pk
 
 
 class SubmissionOut(Schema):