diff --git a/README.md b/README.md index 7c6b064..5831a02 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ LLMs are great at generating options and bad at telling you which one a real per | `list_surveys` | List all your surveys | | `pause_survey` | Pause task serving for an active survey (in-flight responses keep arriving) | | `resume_survey` | Resume task serving for a paused survey | +| `cancel_survey` | Permanently cancel a survey and refund unused reserved credits (irreversible) | | `check_balance` | Check your account balance | | `add_credits` | Open a checkout link to top up your account | diff --git a/mcp_server/client.py b/mcp_server/client.py index dba519a..05c0c08 100644 --- a/mcp_server/client.py +++ b/mcp_server/client.py @@ -122,6 +122,10 @@ def resume_job(self, job_id: str) -> dict: """POST /jobs/{job_id}/resume — resumes task serving.""" return self._request("POST", f"/jobs/{job_id}/resume") + def cancel_job(self, job_id: str) -> dict: + """POST /jobs/{job_id}/cancel — irreversibly cancel and settle reserved credits.""" + return self._request("POST", f"/jobs/{job_id}/cancel") + def plan_survey(self, description: str, preferences: dict | None = None) -> dict: body: dict = {"description": description} if preferences: diff --git a/mcp_server/server.py b/mcp_server/server.py index 418fd5a..835af49 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -711,14 +711,22 @@ def list_surveys() -> str: def _format_lifecycle_response(verb: str, response: dict) -> str: - """Render the response from a pause/resume call.""" + """Render the response from a pause/resume/cancel call. + + ``cost_usd`` is rendered when the backend includes it (cancel returns the + settled cost; pause/resume don't). + """ job_id = response.get("job_id", "?") status = response.get("status", "?") is_paused = response.get("is_paused", False) - return f"{verb} survey {job_id}. Status: {status}, is_paused: {str(is_paused).lower()}." + line = f"{verb} survey {job_id}. Status: {status}, is_paused: {str(is_paused).lower()}." + cost = response.get("cost_usd") + if cost is not None: + line += f" Settled cost: ${cost:.2f}." + return line -_LIFECYCLE_PAST = {"pause": "Paused", "resume": "Resumed"} +_LIFECYCLE_PAST = {"pause": "Paused", "resume": "Resumed", "cancel": "Cancelled"} def _run_lifecycle_action(verb: str, client_method, job_id: str) -> str: @@ -762,6 +770,29 @@ def resume_survey(job_id: str) -> str: return _run_lifecycle_action("resume", client.resume_job, job_id) +@mcp.tool() +def cancel_survey(job_id: str) -> str: + """Permanently cancel a survey and settle its reserved credits. + + ⚠ Cancellation is irreversible. Once cancelled, the survey stops serving + new tasks, any in-flight responses still complete, and all unused + reserved credits are returned to the balance. Already-consumed credit + (for responses already collected) is NOT refunded — the response is shown + in the result so the user knows the final cost. + + Before calling, show the user the survey id and confirm they want to + cancel. Prefer `pause_survey` when the user just wants to stop temporarily. + + Backend rejects with 400 if the survey is already in a terminal state + (completed, failed, or cancelled) — the message will say which. + + Args: + job_id: The job ID returned by create_survey. + """ + client = _get_client() + return _run_lifecycle_action("cancel", client.cancel_job, job_id) + + @mcp.tool() def retry_failed_datapoints(job_id: str, datapoint_indices: list[int] | None = None) -> str: """Re-queue failed datapoints on a survey. diff --git a/tests/test_client.py b/tests/test_client.py index b35bbcb..69ce71c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -52,5 +52,20 @@ def test_both_flags_propagate(self): self.assertTrue(params.get("include_in_progress")) +class CancelJobTests(unittest.TestCase): + def test_cancel_job_posts_to_cancel_path(self): + client = _make_client() + with mock.patch.object(client, "_request", return_value={}) as req: + client.cancel_job("job_x") + req.assert_called_once_with("POST", "/jobs/job_x/cancel") + + def test_cancel_job_returns_request_payload(self): + client = _make_client() + payload = {"job_id": "job_x", "status": "cancelled", "is_paused": True, "cost_usd": 1.23} + with mock.patch.object(client, "_request", return_value=payload): + result = client.cancel_job("job_x") + self.assertEqual(result, payload) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_error_mapping.py b/tests/test_error_mapping.py index 453eca6..ad4a9e5 100644 --- a/tests/test_error_mapping.py +++ b/tests/test_error_mapping.py @@ -13,6 +13,7 @@ from mcp_server.client import DatapointAPIError from mcp_server.server import ( _describe_upload_error, + cancel_survey, check_balance, check_survey, create_survey, @@ -154,6 +155,48 @@ def test_retry_400_surfaces_backend_reason(self): self.assertIn("Cannot retry: Job is not in a retriable state", out) +class CancelSurveyTests(unittest.TestCase): + def test_cancel_success_renders_settled_cost(self): + client = mock.Mock() + client.cancel_job.return_value = { + "job_id": "job_x", + "status": "cancelled", + "is_paused": True, + "cost_usd": 2.50, + } + with mock.patch("mcp_server.server._get_client", return_value=client): + out = cancel_survey("job_x") + client.cancel_job.assert_called_once_with("job_x") + self.assertIn("Cancelled survey job_x", out) + self.assertIn("Status: cancelled", out) + self.assertIn("is_paused: true", out) + self.assertIn("Settled cost: $2.50.", out) + + def test_cancel_404_renders_not_found(self): + client = mock.Mock() + client.cancel_job.side_effect = DatapointAPIError(404, "Job not found") + with mock.patch("mcp_server.server._get_client", return_value=client): + out = cancel_survey("job_x") + self.assertEqual(out, "Survey not found: job_x") + + def test_cancel_400_already_terminal_surfaces_backend_reason(self): + client = mock.Mock() + client.cancel_job.side_effect = DatapointAPIError( + 400, "Cannot cancel a job with status 'completed'." + ) + with mock.patch("mcp_server.server._get_client", return_value=client): + out = cancel_survey("job_x") + self.assertIn("Cannot cancel:", out) + self.assertIn("status 'completed'", out) + + def test_cancel_500_surfaces_detail(self): + client = mock.Mock() + client.cancel_job.side_effect = DatapointAPIError(500, "internal error") + with mock.patch("mcp_server.server._get_client", return_value=client): + out = cancel_survey("job_x") + self.assertIn("Error: internal error", out) + + class CheckSurveyAudienceTargetingTests(unittest.TestCase): def _status(self, **overrides) -> dict: base = { diff --git a/tests/test_renderers.py b/tests/test_renderers.py index a9e069c..b04f478 100644 --- a/tests/test_renderers.py +++ b/tests/test_renderers.py @@ -586,6 +586,27 @@ def test_resumed_response(self): out = _format_lifecycle_response("Resumed", {"job_id": "job_y", "status": "active", "is_paused": False}) self.assertEqual(out, "Resumed survey job_y. Status: active, is_paused: false.") + def test_cancelled_response_includes_settled_cost(self): + out = _format_lifecycle_response( + "Cancelled", + {"job_id": "job_z", "status": "cancelled", "is_paused": True, "cost_usd": 1.23}, + ) + self.assertEqual( + out, + "Cancelled survey job_z. Status: cancelled, is_paused: true. Settled cost: $1.23.", + ) + + def test_cancelled_response_with_zero_cost(self): + out = _format_lifecycle_response( + "Cancelled", + {"job_id": "job_z", "status": "cancelled", "is_paused": True, "cost_usd": 0.0}, + ) + self.assertIn("Settled cost: $0.00.", out) + + def test_pause_response_omits_cost_when_absent(self): + out = _format_lifecycle_response("Paused", {"job_id": "job_x", "status": "active", "is_paused": True}) + self.assertNotIn("Settled cost", out) + class FormatListSurveysTests(unittest.TestCase): def test_empty_list_message(self):