From 0d60f9b246a3c7700b1e4fe2d91619a09a93d69b Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Wed, 24 Jun 2026 14:44:10 +0800
Subject: [PATCH] feat: add kill_job tool

---
 spark_executor/tools/kill.py | 23 +++++++++++++++++++++++
 tests/unit/test_kill_tool.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 spark_executor/tools/kill.py
 create mode 100644 tests/unit/test_kill_tool.py

diff --git a/spark_executor/tools/kill.py b/spark_executor/tools/kill.py
new file mode 100644
index 0000000..da692c1
--- /dev/null
+++ b/spark_executor/tools/kill.py
@@ -0,0 +1,23 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
+from common.logging import logger
+from spark_executor.core.job_store import JobStore
+from spark_executor.core.yarn_client import kill_application
+
+store = JobStore()
+
+
+def kill_job(job_id: str) -> dict[str, str]:
+    job = store.get(job_id)
+    if job is None:
+        raise KeyError(f"Unknown job_id: {job_id}")
+    kill_application(job.application_id)
+    logger.info(f"kill_job job_id={job_id} application_id={job.application_id}")
+    return {
+        "job_id": job_id,
+        "application_id": job.application_id,
+        "status": "KILLED",
+    }
diff --git a/tests/unit/test_kill_tool.py b/tests/unit/test_kill_tool.py
new file mode 100644
index 0000000..bf7e490
--- /dev/null
+++ b/tests/unit/test_kill_tool.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+from datetime import datetime
+from unittest.mock import patch
+
+from spark_executor.core.job_store import JobStore
+from spark_executor.models import Job
+from spark_executor.tools import kill
+
+
+def test_kill_job_calls_yarn_kill():
+    kill.store = JobStore()
+    kill.store.put(
+        Job(
+            job_id="abc",
+            application_id="application_1",
+            script_path="/tmp/j.py",
+            queue="default",
+            submit_time=datetime(2026, 6, 24),
+            connection="prod",
+        )
+    )
+    with patch("spark_executor.tools.kill.kill_application") as m:
+        result = kill.kill_job("abc")
+    m.assert_called_once_with("application_1")
+    assert result == {
+        "job_id": "abc",
+        "application_id": "application_1",
+        "status": "KILLED",
+    }
+
+
+def test_kill_job_raises_for_unknown_job():
+    kill.store = JobStore()
+    import pytest
+    with pytest.raises(KeyError):
+        kill.kill_job("missing")