Support to delete running/pending Experiment

alphatrion/storage/sqlstore.py

@@ -743,6 +743,16 @@ def list_exps_by_timeframe(
         return exps
 
     def delete_experiment(self, experiment_id: uuid.UUID) -> bool:
+        """
+        Delete a single experiment by setting is_del flag.
+        Also deletes all associated runs.
+
+        Status transitions on deletion:
+        - PENDING experiments -> ABORTED
+        - RUNNING experiments -> CANCELLED
+        - Other statuses remain unchanged
+        """
+
         session = self._session()
 
         # Try to delete the experiment
@@ -752,17 +762,19 @@ def delete_experiment(self, experiment_id: uuid.UUID) -> bool:
             .first()
         )
 
-        if exp and exp.status == Status.RUNNING:
-            raise ValueError(
-                "Cannot delete a running experiment. Please stop it first."
-            )
-
         # Delete all runs associated with this experiment
-        # (regardless of experiment status)
         session.query(Run).filter(Run.experiment_id == experiment_id).update(
             {Run.is_del: 1}, synchronize_session=False
         )
+
         if exp:
+            # Update status based on current state
+            if exp.status == Status.PENDING:
+                exp.status = Status.ABORTED
+            elif exp.status == Status.RUNNING:
+                exp.status = Status.CANCELLED
+            # Other statuses remain unchanged
+
             exp.is_del = 1
             session.commit()
             session.close()

alphatrion/tracing/tracing.py

@@ -1,4 +1,3 @@
-
 from opentelemetry.semconv_ai import TraceloopSpanKindValues
 from traceloop.sdk.decorators import agent as _agent
 from traceloop.sdk.decorators import task as _task

tests/integration/server/test_graphql_mutation.py

@@ -897,10 +897,10 @@ def test_delete_experiments_batch_deletes_runs(
     assert metadb.get_run(run_id=run_2_2) is None
 
 
-def test_delete_running_experiment_fails(
+def test_delete_running_experiment_succeeds(
     execute_graphql, test_org_id, test_user_id, test_team_id
 ):
-    """Test that deleting a running experiment raises an error"""
+    """Test that deleting a running experiment marks it as CANCELLED"""
     runtime.init()
     metadb = runtime.storage_runtime().metadb
 
@@ -913,8 +913,6 @@ def test_delete_running_experiment_fails(
 
     # Set experiment status to RUNNING
     metadb.update_experiment(
-        org_id=test_org_id,
-        team_id=test_team_id,
         experiment_id=experiment_id,
         status=Status.RUNNING,
     )
@@ -924,7 +922,7 @@ def test_delete_running_experiment_fails(
     assert exp is not None
     assert exp.status == Status.RUNNING
 
-    # Try to delete running experiment via mutation
+    # Delete running experiment via mutation
     mutation = f"""
     mutation {{
         deleteExperiment(experimentId: "{experiment_id}")
@@ -936,14 +934,22 @@ def test_delete_running_experiment_fails(
         user_id=test_user_id,
     )
 
-    # Should return an error
-    assert response.errors is not None
-    assert "Cannot delete a running experiment" in str(response.errors[0])
+    # Should succeed without errors
+    assert response.errors is None
+    assert response.data["deleteExperiment"] is True
 
-    # Verify experiment still exists
+    # Verify experiment is deleted (get_experiment filters out deleted ones)
     exp = metadb.get_experiment(experiment_id=experiment_id)
-    assert exp is not None
-    assert exp.status == Status.RUNNING
+    assert exp is None
+
+    # Verify status is CANCELLED by querying directly
+    from alphatrion.storage.sql_models import Experiment
+
+    with metadb._session() as session:
+        exp = session.query(Experiment).filter(Experiment.uuid == experiment_id).first()
+        assert exp is not None
+        assert exp.is_del == 1
+        assert exp.status == Status.CANCELLED
 
 
 def test_delete_experiments_with_running_and_pending(

tests/unit/storage/test_sql.py

@@ -510,3 +510,135 @@ def test_delete_experiments_already_deleted(db):
     # Try to delete again
     deleted_count = db.delete_experiments([exp_id])
     assert deleted_count == 0
+
+
+def test_delete_single_experiment_running(db):
+    """Test that delete_experiment marks RUNNING experiment as CANCELLED"""
+    from alphatrion.storage.sql_models import Experiment
+
+    org_id = uuid.uuid4()
+    team_id = uuid.uuid4()
+    user_id = uuid.uuid4()
+
+    # Create running experiment
+    exp_id = db.create_experiment(
+        org_id=org_id,
+        team_id=team_id,
+        user_id=user_id,
+        name="running_exp",
+        status=Status.RUNNING,
+    )
+
+    # Delete the running experiment
+    result = db.delete_experiment(exp_id)
+    assert result is True
+
+    # Verify experiment is deleted
+    exp = db.get_experiment(exp_id)
+    assert exp is None
+
+    # Verify status is CANCELLED by querying directly
+    with db._session() as session:
+        exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first()
+        assert exp is not None
+        assert exp.is_del == 1
+        assert exp.status == Status.CANCELLED
+
+
+def test_delete_single_experiment_pending(db):
+    """Test that delete_experiment marks PENDING experiment as ABORTED"""
+    from alphatrion.storage.sql_models import Experiment
+
+    org_id = uuid.uuid4()
+    team_id = uuid.uuid4()
+    user_id = uuid.uuid4()
+
+    # Create pending experiment
+    exp_id = db.create_experiment(
+        org_id=org_id,
+        team_id=team_id,
+        user_id=user_id,
+        name="pending_exp",
+        status=Status.PENDING,
+    )
+
+    # Delete the pending experiment
+    result = db.delete_experiment(exp_id)
+    assert result is True
+
+    # Verify experiment is deleted
+    exp = db.get_experiment(exp_id)
+    assert exp is None
+
+    # Verify status is ABORTED by querying directly
+    with db._session() as session:
+        exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first()
+        assert exp is not None
+        assert exp.is_del == 1
+        assert exp.status == Status.ABORTED
+
+
+def test_delete_single_experiment_completed(db):
+    """Test that delete_experiment keeps COMPLETED status unchanged"""
+    from alphatrion.storage.sql_models import Experiment
+
+    org_id = uuid.uuid4()
+    team_id = uuid.uuid4()
+    user_id = uuid.uuid4()
+
+    # Create completed experiment
+    exp_id = db.create_experiment(
+        org_id=org_id,
+        team_id=team_id,
+        user_id=user_id,
+        name="completed_exp",
+        status=Status.COMPLETED,
+    )
+
+    # Delete the completed experiment
+    result = db.delete_experiment(exp_id)
+    assert result is True
+
+    # Verify experiment is deleted
+    exp = db.get_experiment(exp_id)
+    assert exp is None
+
+    # Verify status remains COMPLETED
+    with db._session() as session:
+        exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first()
+        assert exp is not None
+        assert exp.is_del == 1
+        assert exp.status == Status.COMPLETED
+
+
+def test_delete_single_experiment_with_runs(db):
+    """Test that delete_experiment also deletes associated runs"""
+    org_id = uuid.uuid4()
+    team_id = uuid.uuid4()
+    user_id = uuid.uuid4()
+
+    # Create experiment with runs
+    exp_id = db.create_experiment(
+        org_id=org_id,
+        team_id=team_id,
+        user_id=user_id,
+        name="exp_with_runs",
+        status=Status.RUNNING,
+    )
+
+    run_id_1 = db.create_run(
+        org_id=org_id, team_id=team_id, user_id=user_id, experiment_id=exp_id
+    )
+    run_id_2 = db.create_run(
+        org_id=org_id, team_id=team_id, user_id=user_id, experiment_id=exp_id
+    )
+
+    # Delete the experiment
+    result = db.delete_experiment(exp_id)
+    assert result is True
+
+    # Verify runs are also deleted
+    run1 = db.get_run(run_id_1)
+    run2 = db.get_run(run_id_2)
+    assert run1 is None
+    assert run2 is None