forked from molecule-ai/molecule-core
test(integration): crash-resume integration tests for Temporal checkpoints (#790)
Closes #790. Depends on feat/issue-583-1-checkpoint-persistence (PR #788). Platform (Go) — checkpoints_integration_test.go (5 new tests): 1. ThreeStepPersistence: POST task_receive/llm_call/task_complete → GET returns all 3 in step_index DESC order with correct names and payloads. 2. CrashResume_HighestStepIsResumptionPoint: POST steps 0+1 only (crash before step 2) → GET shows step_index=1 as the resume point; task_complete absent. 3. UpsertIdempotency_LatestPayloadWins: POST same (wf_id, step_name) twice with different payloads → List returns only the second payload (ON CONFLICT DO UPDATE). 4. PostCascadeDelete_Returns404: simulate post ON-DELETE-CASCADE state (empty rows) → List returns 404 as expected after workspace deletion. 5. AuthGate_NoToken_Returns401: router-level test with WorkspaceAuth middleware; POST/GET/DELETE all return 401 without a bearer token (no DB calls made). workspace-template — _save_checkpoint + 4 Python tests: - Add async _save_checkpoint() to temporal_workflow.py: POST to the platform checkpoint endpoint after each activity stage; fully non-fatal (try/except inside the function, plus defence-in-depth try/except at every call site). - 4 new pytest cases (test_temporal_workflow.py): - nonfatal_on_http_error: _save_checkpoint raises HTTPStatusError (500) → task_receive_activity still returns {"status":"received"}. - nonfatal_on_network_error: _save_checkpoint raises ConnectError → llm_call_activity still returns success LLMResult. - success_path: _save_checkpoint no-op → activity returns correctly; checkpoint called with correct args. - standalone_http_error_is_swallowed: real _save_checkpoint function swallows HTTP 500 from a mocked httpx.AsyncClient; returns None. All 36 temporal workflow Python tests pass. Go tests: Go binary not in this container; test file verified for syntax and against the sqlmock patterns used throughout the handlers package. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Molecule AI QA Engineer
parent
7c4123e6bd
commit
a663c8de81
484
platform/internal/handlers/checkpoints_integration_test.go
Normal file
484
platform/internal/handlers/checkpoints_integration_test.go
Normal file
@ -0,0 +1,484 @@
|
||||
package handlers
|
||||
|
||||
// checkpoints_integration_test.go
|
||||
//
|
||||
// Integration-level tests for the Temporal checkpoint crash-resume system
|
||||
// (issue #790). These scenarios test multi-step lifecycle flows, access
|
||||
// control at the router level, and idempotent upsert semantics — distinct
|
||||
// from checkpoints_test.go which focuses on single-handler correctness.
|
||||
//
|
||||
// All tests use sqlmock + httptest to stay in-process. Cascade-delete
|
||||
// semantics are verified by simulating the post-cascade state (empty rows)
|
||||
// because ON DELETE CASCADE is enforced by the DB schema, not app code.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Molecule-AI/molecule-monorepo/platform/internal/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// checkpointCols is the column list returned by List queries.
|
||||
var checkpointCols = []string{
|
||||
"id", "workspace_id", "workflow_id", "step_name", "step_index",
|
||||
"completed_at", "payload",
|
||||
}
|
||||
|
||||
// upsertSQL is the pattern matched by sqlmock for the checkpoint upsert.
|
||||
const upsertSQL = "INSERT INTO workflow_checkpoints"
|
||||
|
||||
// selectSQL is the pattern matched by sqlmock for the checkpoint list query.
|
||||
const selectSQL = "SELECT id, workspace_id, workflow_id, step_name, step_index"
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 1 — Checkpoint persistence: all three Temporal stages stored & listed
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestCheckpointsIntegration_ThreeStepPersistence verifies the full three-stage
|
||||
// workflow lifecycle: POST task_receive (step 0) → POST llm_call (step 1) →
|
||||
// POST task_complete (step 2) → GET returns all three in step_index DESC order.
|
||||
//
|
||||
// This mirrors what TemporalWorkflowWrapper calls in temporal_workflow.py
|
||||
// after each of the three activity stages.
|
||||
func TestCheckpointsIntegration_ThreeStepPersistence(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
h := newCheckpointsHandler(t, mock)
|
||||
|
||||
stages := []struct {
|
||||
stepName string
|
||||
stepIndex int
|
||||
id string
|
||||
payload string
|
||||
}{
|
||||
{"task_receive", 0, "ckpt-tr", `{"task_id":"t-1"}`},
|
||||
{"llm_call", 1, "ckpt-lc", `{"model":"claude-sonnet-4-5"}`},
|
||||
{"task_complete", 2, "ckpt-tc", `{"success":true}`},
|
||||
}
|
||||
|
||||
// POST all three stages in order.
|
||||
for _, s := range stages {
|
||||
mock.ExpectQuery(upsertSQL).
|
||||
WithArgs("ws-1", "wf-temporal-001", s.stepName, s.stepIndex, s.payload).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(s.id))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-1"}}
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"workflow_id": "wf-temporal-001",
|
||||
"step_name": s.stepName,
|
||||
"step_index": s.stepIndex,
|
||||
"payload": json.RawMessage(s.payload),
|
||||
})
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.Upsert(c)
|
||||
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("stage %q: expected 201, got %d: %s", s.stepName, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// GET — DB returns them in step_index DESC (task_complete first).
|
||||
mock.ExpectQuery(selectSQL).
|
||||
WithArgs("ws-1", "wf-temporal-001").
|
||||
WillReturnRows(sqlmock.NewRows(checkpointCols).
|
||||
AddRow("ckpt-tc", "ws-1", "wf-temporal-001", "task_complete", 2, "2026-04-17T10:02:00Z", []byte(`{"success":true}`)).
|
||||
AddRow("ckpt-lc", "ws-1", "wf-temporal-001", "llm_call", 1, "2026-04-17T10:01:00Z", []byte(`{"model":"claude-sonnet-4-5"}`)).
|
||||
AddRow("ckpt-tr", "ws-1", "wf-temporal-001", "task_receive", 0, "2026-04-17T10:00:00Z", []byte(`{"task_id":"t-1"}`)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{
|
||||
{Key: "id", Value: "ws-1"},
|
||||
{Key: "wfid", Value: "wf-temporal-001"},
|
||||
}
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
h.List(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("List: expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var result []map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatalf("List: invalid JSON response: %v", err)
|
||||
}
|
||||
if len(result) != 3 {
|
||||
t.Fatalf("expected 3 checkpoints, got %d", len(result))
|
||||
}
|
||||
// Verify step_index DESC ordering (highest first).
|
||||
expectedOrder := []string{"task_complete", "llm_call", "task_receive"}
|
||||
for i, want := range expectedOrder {
|
||||
if got := result[i]["step_name"]; got != want {
|
||||
t.Errorf("result[%d].step_name: want %q, got %v", i, want, got)
|
||||
}
|
||||
}
|
||||
// Verify step_index values.
|
||||
for i, wantIdx := range []float64{2, 1, 0} {
|
||||
if got := result[i]["step_index"]; got != wantIdx {
|
||||
t.Errorf("result[%d].step_index: want %.0f, got %v", i, wantIdx, got)
|
||||
}
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 2 — Crash-and-resume: highest persisted step_index is the resume point
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestCheckpointsIntegration_CrashResume_HighestStepIsResumptionPoint simulates
|
||||
// a process crash after llm_call completes (step 1 persisted) but before
|
||||
// task_complete runs (step 2 never persisted).
|
||||
//
|
||||
// On restart, the workflow queries its checkpoints: the highest step_index
|
||||
// present is 1 (llm_call). The workflow can therefore skip task_receive
|
||||
// and llm_call and resume from task_complete, avoiding duplicate LLM calls.
|
||||
func TestCheckpointsIntegration_CrashResume_HighestStepIsResumptionPoint(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
h := newCheckpointsHandler(t, mock)
|
||||
|
||||
// Two stages persisted before crash.
|
||||
for _, stage := range []struct {
|
||||
name string
|
||||
idx int
|
||||
id string
|
||||
}{
|
||||
{"task_receive", 0, "ckpt-tr"},
|
||||
{"llm_call", 1, "ckpt-lc"},
|
||||
} {
|
||||
mock.ExpectQuery(upsertSQL).
|
||||
WithArgs("ws-crash", "wf-crash-001", stage.name, stage.idx, "null").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(stage.id))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-crash"}}
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"workflow_id": "wf-crash-001",
|
||||
"step_name": stage.name,
|
||||
"step_index": stage.idx,
|
||||
})
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
h.Upsert(c)
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("stage %q: expected 201, got %d", stage.name, w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// On restart: query checkpoints — DB returns step_index DESC.
|
||||
mock.ExpectQuery(selectSQL).
|
||||
WithArgs("ws-crash", "wf-crash-001").
|
||||
WillReturnRows(sqlmock.NewRows(checkpointCols).
|
||||
AddRow("ckpt-lc", "ws-crash", "wf-crash-001", "llm_call", 1, "2026-04-17T10:01:00Z", nil).
|
||||
AddRow("ckpt-tr", "ws-crash", "wf-crash-001", "task_receive", 0, "2026-04-17T10:00:00Z", nil))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{
|
||||
{Key: "id", Value: "ws-crash"},
|
||||
{Key: "wfid", Value: "wf-crash-001"},
|
||||
}
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
h.List(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("List after crash: expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var result []map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatalf("invalid JSON: %v", err)
|
||||
}
|
||||
if len(result) != 2 {
|
||||
t.Fatalf("expected 2 checkpoints (crash before step 2), got %d", len(result))
|
||||
}
|
||||
|
||||
// The first element (highest step_index) is the resumption point.
|
||||
resumeStep := result[0]
|
||||
if resumeStep["step_name"] != "llm_call" {
|
||||
t.Errorf("resume point: want step_name 'llm_call', got %v", resumeStep["step_name"])
|
||||
}
|
||||
if resumeStep["step_index"] != float64(1) {
|
||||
t.Errorf("resume point: want step_index 1, got %v", resumeStep["step_index"])
|
||||
}
|
||||
|
||||
// task_complete (step 2) must be absent.
|
||||
for _, cp := range result {
|
||||
if cp["step_name"] == "task_complete" {
|
||||
t.Error("task_complete should not be present — crash happened before that step")
|
||||
}
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 3 — Upsert idempotency: latest payload wins on repeated POST
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestCheckpointsIntegration_UpsertIdempotency_LatestPayloadWins verifies
|
||||
// that POSTing the same (workspace_id, workflow_id, step_name) triple a second
|
||||
// time with a different payload replaces the stored payload (ON CONFLICT DO UPDATE).
|
||||
//
|
||||
// Concrete scenario: llm_call checkpoint is first saved with {"partial":true}
|
||||
// then overwritten with {"partial":false,"tokens":512} when the activity
|
||||
// retries with the full result.
|
||||
func TestCheckpointsIntegration_UpsertIdempotency_LatestPayloadWins(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
h := newCheckpointsHandler(t, mock)
|
||||
|
||||
const wsID = "ws-idem"
|
||||
const wfID = "wf-idem-001"
|
||||
const ckptID = "ckpt-idem"
|
||||
|
||||
// First POST — partial result.
|
||||
firstPayload := `{"partial":true}`
|
||||
mock.ExpectQuery(upsertSQL).
|
||||
WithArgs(wsID, wfID, "llm_call", 1, firstPayload).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(ckptID))
|
||||
|
||||
postCheckpoint(t, h, wsID, wfID, "llm_call", 1, firstPayload)
|
||||
|
||||
// Second POST — full result overwrites via ON CONFLICT DO UPDATE.
|
||||
secondPayload := `{"partial":false,"tokens":512}`
|
||||
mock.ExpectQuery(upsertSQL).
|
||||
WithArgs(wsID, wfID, "llm_call", 1, secondPayload).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(ckptID)) // same ID after update
|
||||
|
||||
postCheckpoint(t, h, wsID, wfID, "llm_call", 1, secondPayload)
|
||||
|
||||
// GET — DB returns a single row with the updated payload.
|
||||
mock.ExpectQuery(selectSQL).
|
||||
WithArgs(wsID, wfID).
|
||||
WillReturnRows(sqlmock.NewRows(checkpointCols).
|
||||
AddRow(ckptID, wsID, wfID, "llm_call", 1, "2026-04-17T10:01:30Z",
|
||||
[]byte(secondPayload)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: wsID}, {Key: "wfid", Value: wfID}}
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
h.List(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("List: expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var result []map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatalf("invalid JSON: %v", err)
|
||||
}
|
||||
if len(result) != 1 {
|
||||
t.Fatalf("expected 1 row (idempotent upsert), got %d", len(result))
|
||||
}
|
||||
|
||||
// The stored payload must reflect the second POST.
|
||||
payloadRaw, _ := json.Marshal(result[0]["payload"])
|
||||
var payloadMap map[string]interface{}
|
||||
json.Unmarshal(payloadRaw, &payloadMap)
|
||||
if payloadMap["partial"] != false {
|
||||
t.Errorf("payload.partial: want false (updated), got %v", payloadMap["partial"])
|
||||
}
|
||||
if payloadMap["tokens"] != float64(512) {
|
||||
t.Errorf("payload.tokens: want 512, got %v", payloadMap["tokens"])
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 4 — Cascade delete: workspace deletion cascades to checkpoints
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestCheckpointsIntegration_PostCascadeDelete_Returns404 verifies the
|
||||
// application's behaviour after ON DELETE CASCADE removes all checkpoint rows
|
||||
// when their parent workspace is deleted.
|
||||
//
|
||||
// The cascade is enforced by the DB schema:
|
||||
// workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE
|
||||
//
|
||||
// This test simulates the post-cascade state: the checkpoints query that runs
|
||||
// after workspace deletion sees an empty result set and returns 404, exactly
|
||||
// as it would if the workspace had never had checkpoints.
|
||||
func TestCheckpointsIntegration_PostCascadeDelete_Returns404(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
h := newCheckpointsHandler(t, mock)
|
||||
|
||||
const wsID = "ws-cascade"
|
||||
const wfID = "wf-cascade-001"
|
||||
|
||||
// Pre-crash: two checkpoints were persisted.
|
||||
for _, stage := range []struct{ name string; idx int; id string }{
|
||||
{"task_receive", 0, "ckpt-tr"},
|
||||
{"llm_call", 1, "ckpt-lc"},
|
||||
} {
|
||||
mock.ExpectQuery(upsertSQL).
|
||||
WithArgs(wsID, wfID, stage.name, stage.idx, "null").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(stage.id))
|
||||
postCheckpointNoPayload(t, h, wsID, wfID, stage.name, stage.idx)
|
||||
}
|
||||
|
||||
// Workspace is deleted (ON DELETE CASCADE fires, checkpoints are gone).
|
||||
// Simulate post-cascade state: List returns empty rows → handler returns 404.
|
||||
mock.ExpectQuery(selectSQL).
|
||||
WithArgs(wsID, wfID).
|
||||
WillReturnRows(sqlmock.NewRows(checkpointCols)) // empty — cascade deleted them
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: wsID}, {Key: "wfid", Value: wfID}}
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
h.List(c)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("post-cascade List: want 404 (no rows), got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 5 — Auth gate: WorkspaceAuth middleware rejects requests without a token
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestCheckpointsIntegration_AuthGate_NoToken_Returns401 tests the checkpoint
|
||||
// endpoints through a full Gin router with the WorkspaceAuth middleware applied.
|
||||
// Every request lacking a valid Authorization: Bearer token must receive 401.
|
||||
//
|
||||
// This pins the security contract established by #351 / Phase 30.1:
|
||||
// no grace period, no fail-open, no existence check before token validation.
|
||||
func TestCheckpointsIntegration_AuthGate_NoToken_Returns401(t *testing.T) {
|
||||
mockDB, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatalf("sqlmock.New: %v", err)
|
||||
}
|
||||
defer mockDB.Close()
|
||||
|
||||
// No DB expectations — strict WorkspaceAuth path short-circuits before
|
||||
// any handler (and therefore before any DB call) when the bearer is absent.
|
||||
|
||||
r := gin.New()
|
||||
wsGroup := r.Group("/workspaces/:id")
|
||||
wsGroup.Use(middleware.WorkspaceAuth(mockDB))
|
||||
{
|
||||
// Handler uses mockDB too; WorkspaceAuth 401s before the handler runs,
|
||||
// so the DB is never queried — any valid *sql.DB pointer works here.
|
||||
cpth := NewCheckpointsHandler(mockDB)
|
||||
wsGroup.POST("/checkpoints", cpth.Upsert)
|
||||
wsGroup.GET("/checkpoints/:wfid", cpth.List)
|
||||
wsGroup.DELETE("/checkpoints/:wfid", cpth.Delete)
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{
|
||||
"POST",
|
||||
"/workspaces/ws-secure/checkpoints",
|
||||
`{"workflow_id":"wf-1","step_name":"task_receive","step_index":0}`,
|
||||
},
|
||||
{
|
||||
"GET",
|
||||
"/workspaces/ws-secure/checkpoints/wf-1",
|
||||
"",
|
||||
},
|
||||
{
|
||||
"DELETE",
|
||||
"/workspaces/ws-secure/checkpoints/wf-1",
|
||||
"",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.method, func(t *testing.T) {
|
||||
var bodyReader *bytes.Reader
|
||||
if tc.body != "" {
|
||||
bodyReader = bytes.NewReader([]byte(tc.body))
|
||||
} else {
|
||||
bodyReader = bytes.NewReader(nil)
|
||||
}
|
||||
|
||||
req, _ := http.NewRequest(tc.method, tc.path, bodyReader)
|
||||
if tc.body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
// Deliberately no Authorization header.
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Errorf("%s %s without token: want 401, got %d: %s",
|
||||
tc.method, tc.path, w.Code, w.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unexpected DB calls during no-token requests: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// postCheckpoint is a test helper that POSTs a checkpoint with a raw JSON
|
||||
// payload string and asserts a 201 response.
|
||||
func postCheckpoint(t *testing.T, h *CheckpointsHandler, wsID, wfID, stepName string, stepIndex int, rawPayload string) {
|
||||
t.Helper()
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: wsID}}
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"workflow_id": wfID,
|
||||
"step_name": stepName,
|
||||
"step_index": stepIndex,
|
||||
"payload": json.RawMessage(rawPayload),
|
||||
})
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
h.Upsert(c)
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("postCheckpoint %q: expected 201, got %d: %s", stepName, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// postCheckpointNoPayload is a test helper that POSTs a checkpoint without
|
||||
// a payload field (stored as JSON null in the DB).
|
||||
func postCheckpointNoPayload(t *testing.T, h *CheckpointsHandler, wsID, wfID, stepName string, stepIndex int) {
|
||||
t.Helper()
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: wsID}}
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"workflow_id": wfID,
|
||||
"step_name": stepName,
|
||||
"step_index": stepIndex,
|
||||
})
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
h.Upsert(c)
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("postCheckpointNoPayload %q: expected 201, got %d: %s", stepName, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
@ -50,6 +50,8 @@ import uuid
|
||||
from datetime import timedelta
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
@ -60,6 +62,72 @@ _TASK_QUEUE = "molecule-agent-tasks"
|
||||
_WORKFLOW_EXECUTION_TIMEOUT = timedelta(minutes=30)
|
||||
_ACTIVITY_START_TO_CLOSE_TIMEOUT = timedelta(minutes=10)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Checkpoint persistence (non-fatal)
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def _save_checkpoint(
|
||||
workspace_id: str,
|
||||
workflow_id: str,
|
||||
step_name: str,
|
||||
step_index: int,
|
||||
payload: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""POST a step checkpoint to the platform.
|
||||
|
||||
Non-fatal: any HTTP error, network failure, or timeout is logged as a
|
||||
WARNING and silently swallowed so the calling activity always continues.
|
||||
Checkpoint loss is survivable; aborting a workflow on a transient DB or
|
||||
network blip is not.
|
||||
|
||||
Args:
|
||||
workspace_id: The workspace whose token is used for auth.
|
||||
workflow_id: Unique ID for this workflow execution (task_id).
|
||||
step_name: Temporal activity stage name
|
||||
(``task_receive`` / ``llm_call`` / ``task_complete``).
|
||||
step_index: 0-based stage index matching the platform schema.
|
||||
payload: Optional JSON-serialisable dict stored as JSONB.
|
||||
|
||||
Reads:
|
||||
PLATFORM_URL Platform base URL (default ``http://localhost:8080``).
|
||||
"""
|
||||
try:
|
||||
from platform_auth import auth_headers as _auth_headers # type: ignore[import]
|
||||
|
||||
platform_url = os.environ.get("PLATFORM_URL", "http://localhost:8080")
|
||||
url = f"{platform_url}/workspaces/{workspace_id}/checkpoints"
|
||||
body: dict = {
|
||||
"workflow_id": workflow_id,
|
||||
"step_name": step_name,
|
||||
"step_index": step_index,
|
||||
}
|
||||
if payload is not None:
|
||||
body["payload"] = payload
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.post(url, json=body, headers=_auth_headers())
|
||||
resp.raise_for_status()
|
||||
|
||||
logger.debug(
|
||||
"Temporal: checkpoint saved workspace=%s wf=%s step=%s idx=%d",
|
||||
workspace_id,
|
||||
workflow_id,
|
||||
step_name,
|
||||
step_index,
|
||||
)
|
||||
except Exception as exc:
|
||||
# Non-fatal: workflow continues regardless of checkpoint outcome.
|
||||
logger.warning(
|
||||
"Temporal: checkpoint failed workspace=%s wf=%s step=%s: %s "
|
||||
"(non-fatal — workflow continues)",
|
||||
workspace_id,
|
||||
workflow_id,
|
||||
step_name,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Serialisable data models
|
||||
# These are the only objects that cross the Temporal serialisation boundary.
|
||||
@ -129,6 +197,9 @@ try:
|
||||
it validates that the in-process registry entry exists and logs receipt.
|
||||
The actual A2A "working" signal (``updater.start_work()``) is emitted
|
||||
inside ``_core_execute()`` so that SSE timing is preserved.
|
||||
|
||||
Saves a step checkpoint after completing. Checkpoint failure is
|
||||
non-fatal — the activity returns normally regardless.
|
||||
"""
|
||||
logger.info(
|
||||
"Temporal[task_receive] task_id=%s context_id=%s workspace=%s model=%s",
|
||||
@ -143,8 +214,22 @@ try:
|
||||
"(crash recovery path — no SSE client connection available)",
|
||||
inp.task_id,
|
||||
)
|
||||
try:
|
||||
await _save_checkpoint(
|
||||
inp.workspace_id, inp.task_id, "task_receive", 0,
|
||||
{"task_id": inp.task_id, "status": "registry_miss"},
|
||||
)
|
||||
except Exception as _ckpt_exc: # pragma: no cover
|
||||
logger.warning("task_receive checkpoint swallowed: %s", _ckpt_exc)
|
||||
return {"task_id": inp.task_id, "status": "registry_miss"}
|
||||
|
||||
try:
|
||||
await _save_checkpoint(
|
||||
inp.workspace_id, inp.task_id, "task_receive", 0,
|
||||
{"task_id": inp.task_id, "status": "received"},
|
||||
)
|
||||
except Exception as _ckpt_exc: # pragma: no cover
|
||||
logger.warning("task_receive checkpoint swallowed: %s", _ckpt_exc)
|
||||
return {"task_id": inp.task_id, "status": "received"}
|
||||
|
||||
@activity.defn(name="llm_call")
|
||||
@ -169,7 +254,15 @@ try:
|
||||
"process likely restarted; original SSE client connection is gone"
|
||||
)
|
||||
logger.warning("Temporal[llm_call] registry miss: %s", msg)
|
||||
return LLMResult(final_text="", success=False, error=msg)
|
||||
miss_result = LLMResult(final_text="", success=False, error=msg)
|
||||
try:
|
||||
await _save_checkpoint(
|
||||
inp.workspace_id, inp.task_id, "llm_call", 1,
|
||||
{"success": False, "error": msg},
|
||||
)
|
||||
except Exception as _ckpt_exc: # pragma: no cover
|
||||
logger.warning("llm_call checkpoint swallowed: %s", _ckpt_exc)
|
||||
return miss_result
|
||||
|
||||
try:
|
||||
executor = entry["executor"]
|
||||
@ -182,7 +275,7 @@ try:
|
||||
|
||||
# Cache for task_complete observability
|
||||
entry["final_text"] = final_text or ""
|
||||
return LLMResult(final_text=final_text or "", success=True)
|
||||
result = LLMResult(final_text=final_text or "", success=True)
|
||||
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
@ -191,7 +284,16 @@ try:
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return LLMResult(final_text="", success=False, error=str(exc))
|
||||
result = LLMResult(final_text="", success=False, error=str(exc))
|
||||
|
||||
try:
|
||||
await _save_checkpoint(
|
||||
inp.workspace_id, inp.task_id, "llm_call", 1,
|
||||
{"success": result.success, "error": result.error or None},
|
||||
)
|
||||
except Exception as _ckpt_exc: # pragma: no cover
|
||||
logger.warning("llm_call checkpoint swallowed: %s", _ckpt_exc)
|
||||
return result
|
||||
|
||||
@activity.defn(name="task_complete")
|
||||
async def task_complete_activity(result: LLMResult) -> None:
|
||||
@ -201,6 +303,11 @@ try:
|
||||
This activity records the outcome for Temporal observability. The actual
|
||||
OTEL task_complete span fires inside ``_core_execute()``; this activity
|
||||
provides a durable, queryable record in Temporal's workflow history.
|
||||
|
||||
Saves a step checkpoint. Checkpoint failure is non-fatal.
|
||||
The ``workspace_id`` and ``task_id`` are not available in this activity
|
||||
(only the ``LLMResult`` is passed from ``llm_call``), so the checkpoint
|
||||
is skipped here — ``llm_call`` already captured the final outcome.
|
||||
"""
|
||||
if result.success:
|
||||
logger.info(
|
||||
|
||||
@ -639,3 +639,242 @@ async def test_molecule_workflow_run_method(real_temporal_with_temporalio):
|
||||
|
||||
assert result is mock_llm_result
|
||||
assert call_count["n"] == 3 # three stages called
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Issue #790 — Case 6: Non-fatal checkpoint failure
|
||||
#
|
||||
# _save_checkpoint() is called from task_receive_activity and llm_call_activity
|
||||
# after their main work completes. If the HTTP POST to the platform returns an
|
||||
# error status (e.g. 500 Internal Server Error) or raises a network exception,
|
||||
# the activity must NOT propagate the error — the workflow continues normally.
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_checkpoint_failure_is_nonfatal_on_http_error(
|
||||
real_temporal_with_temporalio, monkeypatch
|
||||
):
|
||||
"""_save_checkpoint raises httpx.HTTPStatusError (500) → activity succeeds.
|
||||
|
||||
Injects a checkpoint endpoint failure into task_receive_activity by patching
|
||||
_save_checkpoint to raise an HTTPStatusError. The activity must return
|
||||
normally with status='received' regardless.
|
||||
"""
|
||||
mod, _mocks, _mock_shared = real_temporal_with_temporalio
|
||||
|
||||
# Track whether the mock was called.
|
||||
save_calls: list[dict] = []
|
||||
|
||||
async def _fail_checkpoint(workspace_id, workflow_id, step_name, step_index, payload=None):
|
||||
save_calls.append({
|
||||
"workspace_id": workspace_id,
|
||||
"workflow_id": workflow_id,
|
||||
"step_name": step_name,
|
||||
"step_index": step_index,
|
||||
"payload": payload,
|
||||
})
|
||||
# Simulate HTTP 500 from the platform checkpoint endpoint.
|
||||
import httpx as _httpx
|
||||
request = _httpx.Request("POST", "http://localhost:8080/workspaces/ws-1/checkpoints")
|
||||
response = _httpx.Response(500, request=request, text="Internal Server Error")
|
||||
raise _httpx.HTTPStatusError("500", request=request, response=response)
|
||||
|
||||
monkeypatch.setattr(mod, "_save_checkpoint", _fail_checkpoint)
|
||||
|
||||
# Register a minimal task entry so the activity doesn't take the registry-miss path.
|
||||
task_id = "t-nonfatal-ckpt"
|
||||
mod._task_registry[task_id] = {
|
||||
"executor": None,
|
||||
"context": None,
|
||||
"event_queue": None,
|
||||
"final_text": "",
|
||||
}
|
||||
|
||||
inp = mod.AgentTaskInput(
|
||||
task_id=task_id,
|
||||
context_id="ctx-1",
|
||||
user_input="hello",
|
||||
model="test-model",
|
||||
workspace_id="ws-1",
|
||||
history=[],
|
||||
)
|
||||
|
||||
# Act: call task_receive_activity directly. It should succeed despite
|
||||
# _save_checkpoint raising HTTPStatusError.
|
||||
result = await mod.task_receive_activity(inp)
|
||||
|
||||
# Assert: activity returned successfully — checkpoint failure was swallowed.
|
||||
assert result == {"task_id": task_id, "status": "received"}, (
|
||||
f"task_receive_activity must succeed even when checkpoint POST fails; "
|
||||
f"got {result!r}"
|
||||
)
|
||||
# The checkpoint attempt was made (once, for task_receive).
|
||||
assert len(save_calls) == 1
|
||||
assert save_calls[0]["step_name"] == "task_receive"
|
||||
assert save_calls[0]["step_index"] == 0
|
||||
|
||||
# Cleanup registry.
|
||||
mod._task_registry.pop(task_id, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_checkpoint_failure_is_nonfatal_on_network_error(
|
||||
real_temporal_with_temporalio, monkeypatch
|
||||
):
|
||||
"""_save_checkpoint raises a generic network error → llm_call_activity succeeds.
|
||||
|
||||
Tests the llm_call_activity path: even if _save_checkpoint raises a
|
||||
ConnectError (network unreachable), the activity returns its LLMResult.
|
||||
"""
|
||||
mod, _mocks, _mock_shared = real_temporal_with_temporalio
|
||||
|
||||
save_calls: list[str] = []
|
||||
|
||||
async def _network_fail_checkpoint(
|
||||
workspace_id, workflow_id, step_name, step_index, payload=None
|
||||
):
|
||||
save_calls.append(step_name)
|
||||
import httpx as _httpx
|
||||
raise _httpx.ConnectError("Connection refused")
|
||||
|
||||
monkeypatch.setattr(mod, "_save_checkpoint", _network_fail_checkpoint)
|
||||
|
||||
# Build a mock executor whose _core_execute returns a known string.
|
||||
mock_executor = MagicMock()
|
||||
mock_executor._core_execute = AsyncMock(return_value="workflow output")
|
||||
mock_context = MagicMock()
|
||||
mock_event_queue = MagicMock()
|
||||
|
||||
task_id = "t-network-fail"
|
||||
mod._task_registry[task_id] = {
|
||||
"executor": mock_executor,
|
||||
"context": mock_context,
|
||||
"event_queue": mock_event_queue,
|
||||
"final_text": "",
|
||||
}
|
||||
|
||||
inp = mod.AgentTaskInput(
|
||||
task_id=task_id,
|
||||
context_id="ctx-2",
|
||||
user_input="test",
|
||||
model="test-model",
|
||||
workspace_id="ws-2",
|
||||
history=[],
|
||||
)
|
||||
|
||||
# Act: llm_call_activity must complete successfully.
|
||||
result = await mod.llm_call_activity(inp)
|
||||
|
||||
# Assert: successful LLMResult returned despite checkpoint ConnectError.
|
||||
assert isinstance(result, mod.LLMResult), f"Expected LLMResult, got {type(result)}"
|
||||
assert result.success is True, f"llm_call must succeed when checkpoint fails; got {result!r}"
|
||||
assert result.final_text == "workflow output"
|
||||
# _core_execute was called (actual work happened).
|
||||
mock_executor._core_execute.assert_awaited_once_with(mock_context, mock_event_queue)
|
||||
# Checkpoint was attempted (once, for llm_call at step_index=1).
|
||||
assert "llm_call" in save_calls
|
||||
|
||||
mod._task_registry.pop(task_id, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_checkpoint_success_path(
|
||||
real_temporal_with_temporalio, monkeypatch
|
||||
):
|
||||
"""When _save_checkpoint succeeds, activity returns correctly and checkpoint is recorded.
|
||||
|
||||
Verifies the happy path: checkpoint is called with the right arguments and
|
||||
the activity return value is unaffected by a successful checkpoint save.
|
||||
"""
|
||||
mod, _mocks, _mock_shared = real_temporal_with_temporalio
|
||||
|
||||
save_calls: list[dict] = []
|
||||
|
||||
async def _noop_checkpoint(workspace_id, workflow_id, step_name, step_index, payload=None):
|
||||
save_calls.append({
|
||||
"workspace_id": workspace_id,
|
||||
"workflow_id": workflow_id,
|
||||
"step_name": step_name,
|
||||
"step_index": step_index,
|
||||
"payload": payload,
|
||||
})
|
||||
|
||||
monkeypatch.setattr(mod, "_save_checkpoint", _noop_checkpoint)
|
||||
|
||||
task_id = "t-success-ckpt"
|
||||
mod._task_registry[task_id] = {
|
||||
"executor": None,
|
||||
"context": None,
|
||||
"event_queue": None,
|
||||
"final_text": "",
|
||||
}
|
||||
|
||||
inp = mod.AgentTaskInput(
|
||||
task_id=task_id,
|
||||
context_id="ctx-3",
|
||||
user_input="hi",
|
||||
model="test-model",
|
||||
workspace_id="ws-3",
|
||||
history=[],
|
||||
)
|
||||
|
||||
result = await mod.task_receive_activity(inp)
|
||||
|
||||
assert result == {"task_id": task_id, "status": "received"}
|
||||
assert len(save_calls) == 1
|
||||
assert save_calls[0]["workspace_id"] == "ws-3"
|
||||
assert save_calls[0]["workflow_id"] == task_id
|
||||
assert save_calls[0]["step_name"] == "task_receive"
|
||||
assert save_calls[0]["step_index"] == 0
|
||||
|
||||
mod._task_registry.pop(task_id, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_checkpoint_standalone_http_error_is_swallowed(
|
||||
real_temporal_with_temporalio, monkeypatch
|
||||
):
|
||||
"""_save_checkpoint() itself swallows HTTP errors — direct call test.
|
||||
|
||||
Calls the real _save_checkpoint function (patching httpx.AsyncClient)
|
||||
and asserts it returns None without raising even when the platform
|
||||
returns a 500 status.
|
||||
"""
|
||||
import httpx as _httpx
|
||||
|
||||
mod, _mocks, _mock_shared = real_temporal_with_temporalio
|
||||
|
||||
# Patch platform_auth to avoid disk reads in the test environment.
|
||||
mock_platform_auth = MagicMock()
|
||||
mock_platform_auth.auth_headers = MagicMock(return_value={"Authorization": "Bearer test-tok"})
|
||||
monkeypatch.setitem(
|
||||
__import__("sys").modules, "platform_auth", mock_platform_auth
|
||||
)
|
||||
|
||||
# Simulate the AsyncClient.post returning a 500.
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status.side_effect = _httpx.HTTPStatusError(
|
||||
"500",
|
||||
request=_httpx.Request("POST", "http://localhost:8080/workspaces/ws-x/checkpoints"),
|
||||
response=_httpx.Response(500),
|
||||
)
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr(_httpx, "AsyncClient", MagicMock(return_value=mock_client))
|
||||
|
||||
# Must NOT raise — non-fatal contract.
|
||||
result = await mod._save_checkpoint(
|
||||
workspace_id="ws-x",
|
||||
workflow_id="wf-x",
|
||||
step_name="task_receive",
|
||||
step_index=0,
|
||||
payload={"task_id": "t-x"},
|
||||
)
|
||||
|
||||
assert result is None, "_save_checkpoint must return None (no exception) on HTTP 500"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user