diff --git a/cmd/spinifex/cmd/admin_ochre_vector_test.go b/cmd/spinifex/cmd/admin_ochre_vector_test.go index 5d3391018..d0b2221f8 100644 --- a/cmd/spinifex/cmd/admin_ochre_vector_test.go +++ b/cmd/spinifex/cmd/admin_ochre_vector_test.go @@ -30,6 +30,8 @@ type fakeVectorService struct { queryErr error listJobs []handlers_ochrevector.JobRecord listJobsErr error + stopJob handlers_ochrevector.JobRecord + stopJobErr error createIndexInputs []*handlers_ochrevector.CreateIndexRequest deleteIndexInputs []*handlers_ochrevector.DeleteIndexRequest @@ -94,6 +96,17 @@ func (f *fakeVectorService) ListJobs(_ context.Context, _ *handlers_ochrevector. return &handlers_ochrevector.ListJobsResponse{Jobs: f.listJobs}, nil } +func (f *fakeVectorService) StopJob(_ context.Context, req *handlers_ochrevector.StopJobRequest, _ string) (*handlers_ochrevector.StopJobResponse, error) { + if f.stopJobErr != nil { + return nil, f.stopJobErr + } + job := f.stopJob + if job.ID == "" { + job.ID = req.JobID + } + return &handlers_ochrevector.StopJobResponse{Job: job}, nil +} + var _ handlers_ochrevector.VectorService = (*fakeVectorService)(nil) const testIndexID = "idx-0123456789abcdef0" diff --git a/docs/service-interfaces.yaml b/docs/service-interfaces.yaml index 580635ff1..8691caf99 100644 --- a/docs/service-interfaces.yaml +++ b/docs/service-interfaces.yaml @@ -137,6 +137,7 @@ services: - "ochre.vector.describeJob" - "ochre.vector.query" - "ochre.vector.listJobs" + - "ochre.vector.stopJob" publishes: # → vpcd - "vpc.create" diff --git a/spinifex/daemon/ochre_deps.go b/spinifex/daemon/ochre_deps.go index bb317dbe4..6d5599be1 100644 --- a/spinifex/daemon/ochre_deps.go +++ b/spinifex/daemon/ochre_deps.go @@ -161,6 +161,7 @@ func (d *Daemon) startOchreVector() { {handlers_ochrevector.SubjectDescribeJob, handleNATSRequest(vectorService.DescribeJob), "spinifex-workers"}, {handlers_ochrevector.SubjectQuery, handleNATSRequest(vectorService.Query), "spinifex-workers"}, {handlers_ochrevector.SubjectListJobs, handleNATSRequest(vectorService.ListJobs), "spinifex-workers"}, + {handlers_ochrevector.SubjectStopJob, handleNATSRequest(vectorService.StopJob), "spinifex-workers"}, } // backup/restore need RegrantAccount alongside EnsureAccount, which is diff --git a/spinifex/gateway/bedrock/authz.go b/spinifex/gateway/bedrock/authz.go index 1770b93f3..c06df5fc0 100644 --- a/spinifex/gateway/bedrock/authz.go +++ b/spinifex/gateway/bedrock/authz.go @@ -103,6 +103,7 @@ var bedrockScopes = map[string]map[string][]resourceSource{ "StartIngestionJob": {sourceKnowledgeBasePath}, "ListIngestionJobs": {sourceKnowledgeBasePath}, "GetIngestionJob": {sourceKnowledgeBasePath}, + "StopIngestionJob": {sourceKnowledgeBasePath}, }, "bedrock-agent-runtime": { diff --git a/spinifex/gateway/bedrockagent.go b/spinifex/gateway/bedrockagent.go index c04d4e67c..fc43f9374 100644 --- a/spinifex/gateway/bedrockagent.go +++ b/spinifex/gateway/bedrockagent.go @@ -20,6 +20,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/awserrors" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" + "github.com/nats-io/nats.go" ) // bedrockAgentRoute maps one HTTP method + path regex to an AWS action and handler. @@ -100,6 +101,10 @@ var bedrockAgentRoutes = []bedrockAgentRoute{ func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) { return GetIngestionJob(ctx, acct, kb, ds, vector, &bedrockagent.GetIngestionJobInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1]), IngestionJobId: aws.String(p[2])}) }}, + {"PUT", regexp.MustCompile(`^/knowledgebases/([^/]+)/datasources/([^/]+)/ingestionjobs/([^/]+)/stop$`), "StopIngestionJob", + func(ctx context.Context, acct, region string, p []string, b []byte, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService) (any, error) { + return StopIngestionJob(ctx, acct, kb, ds, vector, &StopIngestionJobInput{KnowledgeBaseId: aws.String(p[0]), DataSourceId: aws.String(p[1]), IngestionJobId: aws.String(p[2])}) + }}, } // lookupBedrockAgentAction matches method+path against bedrockAgentRoutes, @@ -247,6 +252,13 @@ func dataSourceStatusToAWS(status string) string { } } +// ingestionJobStatusStopped is StopIngestionJob's terminal wire status. The +// vendored aws-sdk-go v1.55.8 predates AWS's StopIngestionJob operation and +// declares no IngestionJobStatus constant for it; IngestionJob.Status is a +// plain *string, so emitting this value is safe even though the SDK itself +// does not name it. +const ingestionJobStatusStopped = "STOPPED" + // jobStateToAWS translates .9's JobRecord.State to AWS's IngestionJobStatus // wire values. func jobStateToAWS(state string) string { @@ -257,6 +269,8 @@ func jobStateToAWS(state string) string { return bedrockagent.IngestionJobStatusInProgress case handlers_ochrevector.JobStateReady: return bedrockagent.IngestionJobStatusComplete + case handlers_ochrevector.JobStateStopped: + return ingestionJobStatusStopped default: return bedrockagent.IngestionJobStatusFailed } @@ -289,6 +303,8 @@ func translateVectorErr(err error) error { return errors.New(awserrors.ErrorResourceNotFoundException) case errors.Is(err, handlers_ochrevector.ErrIndexExists): return errors.New(awserrors.ErrorConflictException) + case errors.Is(err, nats.ErrNoResponders), errors.Is(err, nats.ErrTimeout), errors.Is(err, context.DeadlineExceeded): + return errors.New(awserrors.ErrorServiceUnavailableException) default: return err } @@ -809,6 +825,70 @@ func GetIngestionJob(ctx context.Context, accountID string, kb *handlers_ochreve return &bedrockagent.GetIngestionJobOutput{IngestionJob: jobRecordToOutput(kbID, dsID, resp.Job)}, nil } +// StopIngestionJobInput/StopIngestionJobOutput mirror the shape of +// GetIngestionJobInput/GetIngestionJobOutput (KnowledgeBaseId/DataSourceId/ +// IngestionJobId in, an IngestionJob summary out). Defined locally, not in +// aws-sdk-go, because the vendored aws-sdk-go v1.55.8 predates AWS's +// StopIngestionJob operation and carries no types for it. +type StopIngestionJobInput struct { + DataSourceId *string `json:"dataSourceId"` + IngestionJobId *string `json:"ingestionJobId"` + KnowledgeBaseId *string `json:"knowledgeBaseId"` +} + +type StopIngestionJobOutput struct { + IngestionJob *bedrockagent.IngestionJob `json:"ingestionJob"` +} + +// StopIngestionJob resolves jobId the same ownership-scoped way +// GetIngestionJob does (both the addressed knowledge base and data source +// must match) before cancelling it via VectorService.StopJob, so a +// foreign/mismatched knowledgeBaseId or dataSourceId in the path cannot be +// used to stop another KB/data source's job by guessing its id. +func StopIngestionJob(ctx context.Context, accountID string, kb *handlers_ochrevector.KBStore, ds *handlers_ochrevector.DataSourceStore, vector handlers_ochrevector.VectorService, input *StopIngestionJobInput) (*StopIngestionJobOutput, error) { + if input == nil || aws.StringValue(input.KnowledgeBaseId) == "" || aws.StringValue(input.DataSourceId) == "" || aws.StringValue(input.IngestionJobId) == "" { + return nil, errors.New(awserrors.ErrorValidationException) + } + kbID := aws.StringValue(input.KnowledgeBaseId) + dsID := aws.StringValue(input.DataSourceId) + jobID := aws.StringValue(input.IngestionJobId) + + kbRec, err := kb.Get(ctx, accountID, kbID) + if err != nil { + return nil, err + } + if kbRec == nil { + return nil, errKBNotFound(kbID) + } + + dsRec, err := ds.Get(ctx, accountID, dsID) + if err != nil { + return nil, err + } + if dsRec == nil || dsRec.KnowledgeBaseID != kbID { + return nil, errDataSourceNotFound(dsID) + } + + // DescribeJob first, purely for the same ownership check GetIngestionJob + // enforces: StopJob alone has no knowledge base/data source to compare + // against, so a foreign job id must be rejected before anything is + // cancelled. + describeResp, err := vector.DescribeJob(ctx, &handlers_ochrevector.DescribeJobRequest{JobID: jobID}, accountID) + if err != nil { + return nil, translateVectorErr(err) + } + if describeResp.Job.IndexID != kbRec.IndexID || describeResp.Job.DataSourceID != dsID { + return nil, errIngestionJobNotFound(jobID) + } + + resp, err := vector.StopJob(ctx, &handlers_ochrevector.StopJobRequest{JobID: jobID}, accountID) + if err != nil { + return nil, translateVectorErr(err) + } + + return &StopIngestionJobOutput{IngestionJob: jobRecordToOutput(kbID, dsID, resp.Job)}, nil +} + // ListIngestionJobs lists knowledgeBaseId's jobs via the new // VectorService.ListJobs, filtered to dataSourceId's own bound index and // exact DataSourceID, and sorted by start time. diff --git a/spinifex/gateway/bedrockagent_request_test.go b/spinifex/gateway/bedrockagent_request_test.go index db9df2c29..d2d8c6085 100644 --- a/spinifex/gateway/bedrockagent_request_test.go +++ b/spinifex/gateway/bedrockagent_request_test.go @@ -304,6 +304,45 @@ func TestBedrockAgentRequest_StartAndListAndGetIngestionJob(t *testing.T) { assert.Equal(t, "job-1", *getOut.IngestionJob.IngestionJobId) } +// TestBedrockAgentRequest_StopIngestionJob proves the StopIngestionJob route +// resolves (lookupBedrockAgentAction no longer falls through to +// ErrInvalidAction for it) and round-trips a STOPPED job through the full +// HTTP dispatch path. +func TestBedrockAgentRequest_StopIngestionJob(t *testing.T) { + vector := &fakeBedrockAgentVectorService{ingestResp: handlers_ochrevector.IngestResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-1", State: handlers_ochrevector.JobStatePending}, + }} + gw := newBedrockAgentRequestGateway(t, vector) + kbID := createTestKnowledgeBase(t, gw) + dsID := createTestDataSource(t, gw, kbID) + + kbRec, err := gw.BedrockAgentKB.Get(context.Background(), bedrockAgentTestAccount, kbID) + require.NoError(t, err) + dsRec, err := gw.BedrockAgentDataSources.Get(context.Background(), bedrockAgentTestAccount, dsID) + require.NoError(t, err) + vector.ingestResp.Job.IndexID = kbRec.IndexID + vector.ingestResp.Job.DataSourceID = dsRec.ID + vector.describeJobResp = handlers_ochrevector.DescribeJobResponse{Job: vector.ingestResp.Job} + stopped := vector.ingestResp.Job + stopped.State = handlers_ochrevector.JobStateStopped + vector.stopJobResp = handlers_ochrevector.StopJobResponse{Job: stopped} + + startReq := bedrockAgentRequestWithAccount(http.MethodPut, "/knowledgebases/"+kbID+"/datasources/"+dsID+"/ingestionjobs/", "") + w := httptest.NewRecorder() + require.NoError(t, gw.BedrockAgent_Request(w, startReq)) + require.Equal(t, http.StatusOK, w.Code) + + stopReq := bedrockAgentRequestWithAccount(http.MethodPut, "/knowledgebases/"+kbID+"/datasources/"+dsID+"/ingestionjobs/job-1/stop", "") + w2 := httptest.NewRecorder() + require.NoError(t, gw.BedrockAgent_Request(w2, stopReq)) + require.Equal(t, http.StatusOK, w2.Code) + + var out StopIngestionJobOutput + require.NoError(t, json.Unmarshal(w2.Body.Bytes(), &out)) + assert.Equal(t, "job-1", *out.IngestionJob.IngestionJobId) + assert.Equal(t, ingestionJobStatusStopped, *out.IngestionJob.Status) +} + func TestBedrockAgentRequest_UnknownRouteReturnsInvalidAction(t *testing.T) { gw := newBedrockAgentRequestGateway(t, &fakeBedrockAgentVectorService{}) diff --git a/spinifex/gateway/bedrockagent_test.go b/spinifex/gateway/bedrockagent_test.go index 148641e92..d20ad1092 100644 --- a/spinifex/gateway/bedrockagent_test.go +++ b/spinifex/gateway/bedrockagent_test.go @@ -9,6 +9,7 @@ package gateway import ( "context" "errors" + "fmt" "log/slog" "strings" "testing" @@ -19,6 +20,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/awserrors" handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/nats-io/nats.go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -48,6 +50,10 @@ type fakeBedrockAgentVectorService struct { queryReq *handlers_ochrevector.QueryRequest queryResp handlers_ochrevector.QueryResponse queryErr error + + stopJobReq *handlers_ochrevector.StopJobRequest + stopJobResp handlers_ochrevector.StopJobResponse + stopJobErr error } func (f *fakeBedrockAgentVectorService) CreateIndex(_ context.Context, req *handlers_ochrevector.CreateIndexRequest, _ string) (*handlers_ochrevector.CreateIndexResponse, error) { @@ -107,6 +113,15 @@ func (f *fakeBedrockAgentVectorService) Query(_ context.Context, req *handlers_o return &resp, nil } +func (f *fakeBedrockAgentVectorService) StopJob(_ context.Context, req *handlers_ochrevector.StopJobRequest, _ string) (*handlers_ochrevector.StopJobResponse, error) { + f.stopJobReq = req + if f.stopJobErr != nil { + return nil, f.stopJobErr + } + resp := f.stopJobResp + return &resp, nil +} + var _ handlers_ochrevector.VectorService = (*fakeBedrockAgentVectorService)(nil) func newBedrockAgentTestStores(t *testing.T) (*handlers_ochrevector.KBStore, *handlers_ochrevector.DataSourceStore) { @@ -157,6 +172,7 @@ func TestJobStateToAWS(t *testing.T) { {handlers_ochrevector.JobStateRunning, bedrockagent.IngestionJobStatusInProgress}, {handlers_ochrevector.JobStateReady, bedrockagent.IngestionJobStatusComplete}, {handlers_ochrevector.JobStateFailed, bedrockagent.IngestionJobStatusFailed}, + {handlers_ochrevector.JobStateStopped, ingestionJobStatusStopped}, } for _, tc := range cases { assert.Equal(t, tc.want, jobStateToAWS(tc.state)) @@ -170,6 +186,37 @@ func TestTranslateVectorErr(t *testing.T) { assert.True(t, awserrors.IsErrorCode(translateVectorErr(handlers_ochrevector.ErrIndexExists), awserrors.ErrorConflictException)) } +// TestTranslateVectorErr_BackendDownMapsToServiceUnavailable proves a +// backend-down error (NATS no-responder/timeout, or a bare context deadline) +// maps to ServiceUnavailableException (503) rather than passing through to +// ErrorHandler's opaque ServerInternal (500) fallback. +func TestTranslateVectorErr_BackendDownMapsToServiceUnavailable(t *testing.T) { + cases := []struct { + name string + err error + }{ + {"no responders", fmt.Errorf("NATS request to ochre.vector.query: %w", nats.ErrNoResponders)}, + {"timeout", fmt.Errorf("NATS request to ochre.vector.query: %w", nats.ErrTimeout)}, + {"deadline exceeded", fmt.Errorf("query: %w", context.DeadlineExceeded)}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := translateVectorErr(tc.err) + require.Error(t, got) + assert.True(t, awserrors.IsErrorCode(got, awserrors.ErrorServiceUnavailableException)) + }) + } +} + +// TestTranslateVectorErr_UnmappedErrorPassesThrough proves an error not +// otherwise classified above passes through unchanged, unaffected by the +// new nats/timeout cases -- ErrorHandler's own fallback still sanitizes it +// to ServerInternal. +func TestTranslateVectorErr_UnmappedErrorPassesThrough(t *testing.T) { + original := errors.New("some other internal failure") + assert.Equal(t, original, translateVectorErr(original)) +} + func validCreateKBInput() *bedrockagent.CreateKnowledgeBaseInput { return &bedrockagent.CreateKnowledgeBaseInput{ Name: aws.String("docs"), @@ -528,3 +575,61 @@ func TestListIngestionJobs_FiltersToBoundIndexAndExactDataSourceID(t *testing.T) assert.Equal(t, "job-match", aws.StringValue(out.IngestionJobSummaries[0].IngestionJobId)) assert.Equal(t, bedrockagent.IngestionJobStatusComplete, aws.StringValue(out.IngestionJobSummaries[0].Status)) } + +// TestStopIngestionJob_CancelsBoundJobAndReturnsStoppedStatus proves a +// StopIngestionJob call resolved against the right knowledge base/data +// source reaches VectorService.StopJob and renders its STOPPED result. +func TestStopIngestionJob_CancelsBoundJobAndReturnsStoppedStatus(t *testing.T) { + kb, ds := newBedrockAgentTestStores(t) + ctx := context.Background() + require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"})) + require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"})) + + vector := &fakeBedrockAgentVectorService{ + describeJobResp: handlers_ochrevector.DescribeJobResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStateRunning}, + }, + stopJobResp: handlers_ochrevector.StopJobResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-1", DataSourceID: "ds-1", State: handlers_ochrevector.JobStateStopped}, + }, + } + out, err := StopIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &StopIngestionJobInput{ + KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), IngestionJobId: aws.String("job-1"), + }) + require.NoError(t, err) + require.NotNil(t, vector.stopJobReq) + assert.Equal(t, "job-1", vector.stopJobReq.JobID) + assert.Equal(t, ingestionJobStatusStopped, aws.StringValue(out.IngestionJob.Status)) +} + +// TestStopIngestionJob_RejectsJobFromForeignIndex mirrors +// TestGetIngestionJob_RejectsJobFromForeignIndex: a job whose IndexID does +// not match the addressed knowledge base's bound index must not be +// cancellable through it. +func TestStopIngestionJob_RejectsJobFromForeignIndex(t *testing.T) { + kb, ds := newBedrockAgentTestStores(t) + ctx := context.Background() + require.NoError(t, kb.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.KBRecord{ID: "kb-1", IndexID: "idx-1"})) + require.NoError(t, ds.Create(ctx, bedrockAgentTestAccount, handlers_ochrevector.DataSourceRecord{ID: "ds-1", KnowledgeBaseID: "kb-1"})) + + vector := &fakeBedrockAgentVectorService{describeJobResp: handlers_ochrevector.DescribeJobResponse{ + Job: handlers_ochrevector.JobRecord{ID: "job-1", IndexID: "idx-other", DataSourceID: "ds-1"}, + }} + _, err := StopIngestionJob(ctx, bedrockAgentTestAccount, kb, ds, vector, &StopIngestionJobInput{ + KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), IngestionJobId: aws.String("job-1"), + }) + require.Error(t, err) + assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException)) + assert.Nil(t, vector.stopJobReq, "StopJob must never be reached for a job that fails the ownership check") +} + +// TestStopIngestionJob_MissingIngestionJobIDIsValidationException proves the +// same input guard GetIngestionJob has applies to StopIngestionJob. +func TestStopIngestionJob_MissingIngestionJobIDIsValidationException(t *testing.T) { + kb, ds := newBedrockAgentTestStores(t) + _, err := StopIngestionJob(context.Background(), bedrockAgentTestAccount, kb, ds, &fakeBedrockAgentVectorService{}, &StopIngestionJobInput{ + KnowledgeBaseId: aws.String("kb-1"), DataSourceId: aws.String("ds-1"), + }) + require.Error(t, err) + assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorValidationException)) +} diff --git a/spinifex/handlers/ochrevector/api.go b/spinifex/handlers/ochrevector/api.go index dd24bf218..d21727caa 100644 --- a/spinifex/handlers/ochrevector/api.go +++ b/spinifex/handlers/ochrevector/api.go @@ -20,6 +20,7 @@ const ( SubjectDescribeJob = "ochre.vector.describeJob" SubjectQuery = "ochre.vector.query" SubjectListJobs = "ochre.vector.listJobs" + SubjectStopJob = "ochre.vector.stopJob" ) // CreateIndexRequest names a new index and its fixed, per-index properties @@ -103,11 +104,21 @@ type ListJobsResponse struct { Jobs []JobRecord `json:"jobs"` } +// StopJobRequest asks to cancel one in-flight ingestion job. +type StopJobRequest struct { + JobID string `json:"jobId"` +} + +type StopJobResponse struct { + Job JobRecord `json:"job"` +} + // VectorService is the tenant vector-store surface: index CRUD, ingestion, -// job status/listing, and similarity query. Every method's accountID argument -// is supplied by the transport alone (a NATS message header for the daemon -// subscription, D10) -- no request type above carries an account field, so a -// spoofed payload account can never widen a caller's own scope. +// job status/listing/cancellation, and similarity query. Every method's +// accountID argument is supplied by the transport alone (a NATS message +// header for the daemon subscription, D10) -- no request type above carries +// an account field, so a spoofed payload account can never widen a caller's +// own scope. type VectorService interface { CreateIndex(ctx context.Context, req *CreateIndexRequest, accountID string) (*CreateIndexResponse, error) DeleteIndex(ctx context.Context, req *DeleteIndexRequest, accountID string) (*DeleteIndexResponse, error) @@ -116,6 +127,7 @@ type VectorService interface { DescribeJob(ctx context.Context, req *DescribeJobRequest, accountID string) (*DescribeJobResponse, error) ListJobs(ctx context.Context, req *ListJobsRequest, accountID string) (*ListJobsResponse, error) Query(ctx context.Context, req *QueryRequest, accountID string) (*QueryResponse, error) + StopJob(ctx context.Context, req *StopJobRequest, accountID string) (*StopJobResponse, error) } // indexService is the CreateIndex/DeleteIndex/ListIndexes surface vectorService @@ -128,10 +140,11 @@ type indexService interface { var _ indexService = (*Service)(nil) -// ingestStarter is the StartIngest surface vectorService delegates to; -// *IngestService satisfies it structurally. +// ingestStarter is the StartIngest/StopJob surface vectorService delegates +// to; *IngestService satisfies it structurally. type ingestStarter interface { StartIngest(ctx context.Context, accountID, indexID string, source SourceSpec, dataSourceID string) (*JobRecord, error) + StopJob(ctx context.Context, accountID, jobID string) (*JobRecord, error) } var _ ingestStarter = (*IngestService)(nil) @@ -284,6 +297,18 @@ func (s *vectorService) ListJobs(ctx context.Context, _ *ListJobsRequest, accoun return &ListJobsResponse{Jobs: jobs}, nil } +// StopJob cancels req.JobID's in-flight run via ingest.StopJob. +func (s *vectorService) StopJob(ctx context.Context, req *StopJobRequest, accountID string) (*StopJobResponse, error) { + if req == nil || req.JobID == "" { + return nil, fmt.Errorf("ochrevector: stop job requires a jobId") + } + job, err := s.ingest.StopJob(ctx, accountID, req.JobID) + if err != nil { + return nil, fmt.Errorf("ochrevector: stop job %s: %w", req.JobID, err) + } + return &StopJobResponse{Job: *job}, nil +} + // Query resolves IndexID's pinned embedding model from the registry (D8), // embeds Text against it, then runs the similarity search directly against // backend, optionally over-fetching and reranking (rerankTopK). A non-READY @@ -376,3 +401,7 @@ func (c *NATSVectorService) ListJobs(ctx context.Context, req *ListJobsRequest, func (c *NATSVectorService) Query(ctx context.Context, req *QueryRequest, accountID string) (*QueryResponse, error) { return utils.NATSRequest[QueryResponse](ctx, c.nc, SubjectQuery, req, c.timeout, accountID) } + +func (c *NATSVectorService) StopJob(ctx context.Context, req *StopJobRequest, accountID string) (*StopJobResponse, error) { + return utils.NATSRequest[StopJobResponse](ctx, c.nc, SubjectStopJob, req, c.timeout, accountID) +} diff --git a/spinifex/handlers/ochrevector/api_test.go b/spinifex/handlers/ochrevector/api_test.go index a6951d95d..9c076d687 100644 --- a/spinifex/handlers/ochrevector/api_test.go +++ b/spinifex/handlers/ochrevector/api_test.go @@ -346,6 +346,7 @@ func subscribeVectorService(t *testing.T, nc *nats.Conn, svc VectorService) { {SubjectDescribeJob, stubVectorNATSHandler(svc.DescribeJob)}, {SubjectQuery, stubVectorNATSHandler(svc.Query)}, {SubjectListJobs, stubVectorNATSHandler(svc.ListJobs)}, + {SubjectStopJob, stubVectorNATSHandler(svc.StopJob)}, } for _, s := range subs { sub, err := nc.QueueSubscribe(s.subject, "spinifex-workers", s.handler) @@ -390,6 +391,12 @@ func TestNATSVectorService_RoundTrip(t *testing.T) { require.Len(t, listJobsOut.Jobs, 1) assert.Equal(t, ingestOut.Job.ID, listJobsOut.Jobs[0].ID) + // The job is PENDING (Sweep has not run in this test), so StopJob finds + // no registered cancel func and reports it back unchanged. + stopOut, err := client.StopJob(ctx, &StopJobRequest{JobID: ingestOut.Job.ID}, apiAccountA) + require.NoError(t, err) + assert.Equal(t, ingestOut.Job.ID, stopOut.Job.ID) + backend.queryResults = []QueryResult{{Chunk: "hi", SourceKey: "docs/a.txt", Score: 0.9}} queryOut, err := client.Query(ctx, &QueryRequest{IndexID: "idx-one", Text: "hello"}, apiAccountA) require.NoError(t, err) diff --git a/spinifex/handlers/ochrevector/ingest.go b/spinifex/handlers/ochrevector/ingest.go index 7aa36dc91..dbe057163 100644 --- a/spinifex/handlers/ochrevector/ingest.go +++ b/spinifex/handlers/ochrevector/ingest.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "log/slog" + "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -60,6 +61,13 @@ const ( // const, so tests can shrink it rather than paying the real delay. var ingestRetryBackoff = 200 * time.Millisecond +// ingestPerDocTimeout bounds one document's ingestObject call (embed calls +// plus both ReplaceDocument calls) against a dead backend, so a document +// fails fast rather than hanging on the caller's root ctx indefinitely -- a +// var, not a const, so tests can shrink it rather than paying the real delay, +// mirroring ingestRetryBackoff. +var ingestPerDocTimeout = 30 * time.Second + // IngestService orchestrates the ingestion job lifecycle: claiming a job, // listing a source bucket/prefix, chunking and embedding each object, and // replacing its rows in the vector backend. @@ -69,11 +77,38 @@ type IngestService struct { Backend VectorBackend Store objectstore.ObjectStore Embedder Embedder + + // cancelMu guards cancelFuncs, the live-run cancel registry StopJob + // consults. + cancelMu sync.Mutex + cancelFuncs map[string]context.CancelFunc } // NewIngestService constructs an IngestService over its dependencies. func NewIngestService(jobs *JobStore, registry *Registry, backend VectorBackend, store objectstore.ObjectStore, embedder Embedder) *IngestService { - return &IngestService{Jobs: jobs, Registry: registry, Backend: backend, Store: store, Embedder: embedder} + return &IngestService{ + Jobs: jobs, Registry: registry, Backend: backend, Store: store, Embedder: embedder, + cancelFuncs: map[string]context.CancelFunc{}, + } +} + +// registerRunCancel records cancel as the way to stop job's in-flight run, +// so a concurrent StopJob call can find it. Overwrites any previous entry +// for the same key -- a job never has two runs in flight at once (Sweep is +// the only caller, and it runs jobs one at a time). +func (s *IngestService) registerRunCancel(accountID, jobID string, cancel context.CancelFunc) { + s.cancelMu.Lock() + defer s.cancelMu.Unlock() + s.cancelFuncs[jobKey(accountID, jobID)] = cancel +} + +// deregisterRunCancel removes job's cancel entry once its run has returned, +// so a later StopJob call for the same (now-finished) job id finds nothing +// to cancel rather than a stale func for a run that already ended. +func (s *IngestService) deregisterRunCancel(accountID, jobID string) { + s.cancelMu.Lock() + defer s.cancelMu.Unlock() + delete(s.cancelFuncs, jobKey(accountID, jobID)) } // StartIngest validates indexID exists for accountID, then reserves a new @@ -118,18 +153,39 @@ func (s *IngestService) ReconcileFromSource(ctx context.Context, accountID, inde // RunJob is the synchronous ingestion worker: lists job.Source, and for each // object, reads/chunks/embeds/replaces its rows. A per-document failure -// (unreadable object, oversize, exhausted embed retries) is skipped and -// recorded rather than failing the whole job; a fully-unavailable embedder -// (every attempted document's embed calls exhaust retries) fails the job -// instead of reaching READY with everything skipped (D7). On success, the -// job's source spec is recorded on the index registry record (D4). +// (unreadable object, oversize, exhausted embed retries, a per-document +// timeout against a dead backend) is skipped and recorded rather than +// failing the whole job; a fully-unavailable embedder (every attempted +// document's embed calls exhaust retries) fails the job instead of reaching +// READY with everything skipped (D7). Progress (DocumentsDone/ +// FailedDocuments) is persisted after every document, so a job on a dead +// backend is visibly making (or not making) progress rather than sitting +// silently IN_PROGRESS. On success, the job's source spec is recorded on the +// index registry record (D4). +// +// runCtx, derived from ctx, is registered so a concurrent StopJob call can +// cancel exactly this run: cancelling it stops the job at JobStateStopped, +// distinct from a per-document timeout (a child of runCtx, scoped to one +// document) reaching JobStateFailed the ordinary per-document way. Every +// JobStore write in this method uses ctx, not runCtx, so the terminal write +// recording a stop is never itself blocked by the cancellation it records. func (s *IngestService) RunJob(ctx context.Context, job JobRecord) error { if err := s.Jobs.SetState(ctx, job.AccountID, job.ID, JobStateRunning); err != nil { return err } - keys, err := s.listObjectKeys(ctx, job.Source) + runCtx, cancel := context.WithCancel(ctx) + s.registerRunCancel(job.AccountID, job.ID, cancel) + defer func() { + s.deregisterRunCancel(job.AccountID, job.ID) + cancel() + }() + + keys, err := s.listObjectKeys(runCtx, job.Source) if err != nil { + if runCtx.Err() != nil { + return s.stopJobRecord(ctx, job, 0, 0, nil) + } return s.failJob(ctx, job, fmt.Errorf("list source objects: %w", err)) } @@ -138,16 +194,37 @@ func (s *IngestService) RunJob(ctx context.Context, job JobRecord) error { embedFailures := 0 for _, key := range keys { - docErr := s.ingestObject(ctx, job.AccountID, job.IndexID, key, job.Source) + if runCtx.Err() != nil { + return s.stopJobRecord(ctx, job, done, len(keys), failedDocs) + } + + docCtx, docCancel := context.WithTimeout(runCtx, ingestPerDocTimeout) + docErr := s.ingestObject(docCtx, job.AccountID, job.IndexID, key, job.Source) + docCancel() + + if docErr != nil && runCtx.Err() != nil { + // A Stop was requested mid-document: the failure is cancellation + // noise, not a real per-document failure to record. + return s.stopJobRecord(ctx, job, done, len(keys), failedDocs) + } + if docErr != nil { slog.WarnContext(ctx, "ochrevector: ingest document failed", "job", job.ID, "key", key, "err", docErr.err) failedDocs = append(failedDocs, FailedDoc{SourceKey: key, Reason: docErr.Error()}) if docErr.fromEmbedder { embedFailures++ } - continue + } else { + done++ + } + + if updErr := s.Jobs.Update(ctx, job.AccountID, job.ID, func(rec *JobRecord) { + rec.DocumentsTotal = len(keys) + rec.DocumentsDone = done + rec.FailedDocuments = failedDocs + }); updErr != nil { + slog.WarnContext(ctx, "ochrevector: persist mid-run ingest progress failed", "job", job.ID, "key", key, "err", updErr) } - done++ } // A job that ingests nothing must not report READY, regardless of why: @@ -174,6 +251,44 @@ func (s *IngestService) RunJob(ctx context.Context, job JobRecord) error { }) } +// stopJobRecord transitions job to JobStateStopped with whatever progress it +// made before its run was cancelled. +func (s *IngestService) stopJobRecord(ctx context.Context, job JobRecord, done, total int, failedDocs []FailedDoc) error { + return s.Jobs.Update(ctx, job.AccountID, job.ID, func(rec *JobRecord) { + rec.State = JobStateStopped + rec.DocumentsTotal = total + rec.DocumentsDone = done + rec.FailedDocuments = failedDocs + }) +} + +// StopJob cancels jobID's in-flight run, if one is currently registered +// (RunJob registers its cancel func for exactly as long as it is running), +// and returns the job's record as of the call -- the run itself finishes the +// transition to JobStateStopped asynchronously, from wherever it currently +// is in RunJob's loop. A job with no registered cancel func -- never +// started, or already finished by the time Stop was requested -- is +// reported back unchanged rather than as an error, so a caller racing the +// job's own natural completion still gets a sane answer. +func (s *IngestService) StopJob(ctx context.Context, accountID, jobID string) (*JobRecord, error) { + job, err := s.Jobs.Get(ctx, accountID, jobID) + if err != nil { + return nil, err + } + if job == nil { + return nil, ErrJobNotFound + } + + s.cancelMu.Lock() + cancel, ok := s.cancelFuncs[jobKey(accountID, jobID)] + s.cancelMu.Unlock() + if ok { + cancel() + } + + return job, nil +} + // failJob transitions job to FAILED with cause recorded, and returns cause // (joined with any error updating the record itself) so the caller sees why. func (s *IngestService) failJob(ctx context.Context, job JobRecord, cause error) error { diff --git a/spinifex/handlers/ochrevector/ingest_test.go b/spinifex/handlers/ochrevector/ingest_test.go index a339f797b..41ed21dda 100644 --- a/spinifex/handlers/ochrevector/ingest_test.go +++ b/spinifex/handlers/ochrevector/ingest_test.go @@ -37,6 +37,11 @@ type stubEmbedder struct { calls int failAll bool failIfContains string + // beforeEmbed, if set, is called with the 1-based call number at the + // start of every Embed call, before failAll/failIfContains are + // evaluated -- lets a test snapshot job progress exactly between two + // documents' embed calls, to prove mid-run persistence. + beforeEmbed func(callNum int) } var _ Embedder = (*stubEmbedder)(nil) @@ -44,8 +49,13 @@ var _ Embedder = (*stubEmbedder)(nil) func (e *stubEmbedder) Embed(_ context.Context, _ string, inputs []string) ([][]float32, error) { e.mu.Lock() e.calls++ + n := e.calls e.mu.Unlock() + if e.beforeEmbed != nil { + e.beforeEmbed(n) + } + if e.failAll { return nil, errors.New("stub embedder: unavailable") } @@ -133,6 +143,35 @@ func (e *tokenCounterStubEmbedder) callCountTokens() int { return e.countCalls } +// blockingEmbedder is a controllable Embedder test double for a fully dead +// backend: every Embed call blocks until its ctx is done (cancelled or +// timed out) and then returns ctx.Err(), the way a real HTTP client would +// hang against an unreachable endpoint absent a bounded context. beforeEmbed, +// if set, is called with the 1-based call number as each call is entered +// (before blocking), so a test can synchronize on "the Nth document's embed +// call is now in flight". +type blockingEmbedder struct { + mu sync.Mutex + calls int + beforeEmbed func(callNum int) +} + +var _ Embedder = (*blockingEmbedder)(nil) + +func (e *blockingEmbedder) Embed(ctx context.Context, _ string, _ []string) ([][]float32, error) { + e.mu.Lock() + e.calls++ + n := e.calls + e.mu.Unlock() + + if e.beforeEmbed != nil { + e.beforeEmbed(n) + } + + <-ctx.Done() + return nil, ctx.Err() +} + // newIngestTestSetupWithEmbedder mirrors newIngestTestSetup but takes a // caller-supplied Embedder, so tests can exercise the TokenLimiter/ // TokenCounter optional-interface wiring without a live TEI endpoint. @@ -723,3 +762,156 @@ func TestRunJob_ConsultsTokenCounterWhenEmbedderProvidesOne(t *testing.T) { assert.Positive(t, embedder.callCountTokens(), "ingestObject must consult the embedder's TokenCounter when it implements one") } + +// TestRunJob_PersistsIncrementalProgressAcrossSuccessfulDocuments proves +// DocumentsDone is written to the job store after every document, not only +// once at the end: the 2nd document's embed call must observe the 1st +// document's completion already persisted. +func TestRunJob_PersistsIncrementalProgressAcrossSuccessfulDocuments(t *testing.T) { + svc, _, _, store, embedder := newIngestTestSetup(t) + ctx := context.Background() + + putObject(t, store, ingestPrefix+"doc1.txt", "first document") + putObject(t, store, ingestPrefix+"doc2.txt", "second document") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource(), "") + require.NoError(t, err) + + docsDoneAtSecondCall := -1 + embedder.beforeEmbed = func(n int) { + if n == 2 { + got, gErr := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, gErr) + docsDoneAtSecondCall = got.DocumentsDone + } + } + + require.NoError(t, svc.RunJob(ctx, *job)) + + assert.Equal(t, 1, docsDoneAtSecondCall, "the 2nd document's embed call must observe the 1st document's progress already persisted, not a single terminal write") +} + +// TestRunJob_DeadBackendFailsPromptlyWithVisibleMidRunProgress proves a job +// against a fully unreachable backend (every embed call hangs past +// ingestPerDocTimeout) reaches FAILED promptly rather than hanging on the +// caller's root ctx, and that FailedDocuments/DocumentsTotal were persisted +// incrementally as each document was attempted, not only in a single +// terminal write -- the 2nd document's embed call must already observe the +// 1st document's failure. +func TestRunJob_DeadBackendFailsPromptlyWithVisibleMidRunProgress(t *testing.T) { + embedder := &blockingEmbedder{} + svc, _, backend, store := newIngestTestSetupWithEmbedder(t, embedder) + ctx := context.Background() + + origTimeout := ingestPerDocTimeout + ingestPerDocTimeout = 20 * time.Millisecond + t.Cleanup(func() { ingestPerDocTimeout = origTimeout }) + + putObject(t, store, ingestPrefix+"doc1.txt", "first document, dead backend") + putObject(t, store, ingestPrefix+"doc2.txt", "second document, dead backend") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource(), "") + require.NoError(t, err) + + var failedDocsAtSecondCall int + var totalAtSecondCall int + embedder.beforeEmbed = func(n int) { + if n == 2 { + got, gErr := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, gErr) + failedDocsAtSecondCall = len(got.FailedDocuments) + totalAtSecondCall = got.DocumentsTotal + } + } + + start := time.Now() + err = svc.RunJob(ctx, *job) + elapsed := time.Since(start) + require.Error(t, err, "a fully dead backend must fail the job, not reach READY with everything skipped") + assert.Less(t, elapsed, 5*time.Second, "a per-document timeout must fail fast rather than hang on the root ctx") + + assert.Equal(t, 1, failedDocsAtSecondCall, "the 2nd document's embed call must observe the 1st document's failure already persisted") + assert.Equal(t, 2, totalAtSecondCall) + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateFailed, got.State) + assert.NotEmpty(t, got.Error) + assert.Equal(t, 0, got.DocumentsDone) + require.Len(t, got.FailedDocuments, 2) + + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", ingestPrefix+"doc1.txt")) +} + +// TestIngestService_StopJob_CancelsRunningJobToStopped proves StopJob +// cancels a job blocked mid-run (on a dead-backend embed call) and the run +// itself transitions to JobStateStopped, not JobStateFailed -- distinct from +// a per-document timeout, which reaches FAILED on its own. +func TestIngestService_StopJob_CancelsRunningJobToStopped(t *testing.T) { + embedder := &blockingEmbedder{} + svc, _, backend, store := newIngestTestSetupWithEmbedder(t, embedder) + ctx := context.Background() + + putObject(t, store, ingestPrefix+"doc1.txt", "a document that will be stopped mid-flight") + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource(), "") + require.NoError(t, err) + + started := make(chan struct{}) + embedder.beforeEmbed = func(int) { close(started) } + + runErr := make(chan error, 1) + go func() { runErr <- svc.RunJob(ctx, *job) }() + + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("RunJob never reached the embedder call") + } + + stopped, err := svc.StopJob(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, stopped) + + select { + case err := <-runErr: + require.NoError(t, err, "a stopped run must not itself return an error") + case <-time.After(5 * time.Second): + t.Fatal("RunJob never returned after StopJob cancelled it") + } + + got, err := svc.Jobs.Get(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStateStopped, got.State) + assert.Equal(t, 0, backend.replaceDocumentCallCount(ingestAccountA, "idx-one", ingestPrefix+"doc1.txt")) +} + +// TestIngestService_StopJob_UnregisteredJobReturnsRecordUnchanged proves a +// job with no registered cancel func (never started running, or already +// finished) is reported back as-is rather than as an error, so a caller +// racing the job's own natural completion still gets a sane answer. +func TestIngestService_StopJob_UnregisteredJobReturnsRecordUnchanged(t *testing.T) { + svc, _, _, _, _ := newIngestTestSetup(t) + ctx := context.Background() + + job, err := svc.StartIngest(ctx, ingestAccountA, "idx-one", testSource(), "") + require.NoError(t, err) + + got, err := svc.StopJob(ctx, ingestAccountA, job.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, JobStatePending, got.State) +} + +// TestIngestService_StopJob_UnknownJobErrors proves StopJob reports +// ErrJobNotFound for a job id that does not exist, mirroring DescribeJob. +func TestIngestService_StopJob_UnknownJobErrors(t *testing.T) { + svc, _, _, _, _ := newIngestTestSetup(t) + ctx := context.Background() + + _, err := svc.StopJob(ctx, ingestAccountA, "job-does-not-exist") + require.Error(t, err) + assert.ErrorIs(t, err, ErrJobNotFound) +} diff --git a/spinifex/handlers/ochrevector/jobs.go b/spinifex/handlers/ochrevector/jobs.go index 884984f5f..ac9fb0e9e 100644 --- a/spinifex/handlers/ochrevector/jobs.go +++ b/spinifex/handlers/ochrevector/jobs.go @@ -24,6 +24,10 @@ const ( JobStateRunning = "RUNNING" JobStateReady = "READY" JobStateFailed = "FAILED" + // JobStateStopped is a terminal state reached only via an explicit + // StopJob call cancelling a RUNNING job -- distinct from JobStateFailed, + // which a per-document timeout or embedder outage reaches on its own. + JobStateStopped = "STOPPED" ) // ErrJobExists reports that Reserve lost the single-writer claim on a job id