feat(kb): add abort capability and enum standardization for zero-downtime update (#274)

Because

- Knowledge Base updates lack graceful cancellation - once started,
updates must complete or fail, with no way to manually abort them when
issues arise
- Status values were represented as untyped strings across the system,
causing inconsistencies between protobuf definitions, backend code,
database storage, and test assertions
- Files that haven't started processing (`NOTSTARTED`, `WAITING`) were
incorrectly blocking KB synchronization, causing integration tests to
timeout indefinitely
- The system lacked comprehensive integration test coverage for update
lifecycle edge cases and concurrent operations during updates

This commit

**Adds Abort API for Knowledge Base Updates**
- **Introduces `AbortKnowledgeBaseUpdateAdmin` API** enabling
administrators to gracefully cancel ongoing KB updates
- **Cancels active Temporal workflows** and cleans up staging KB
resources (both completed and incomplete files)
- **Supports selective abort** - can abort specific catalogs by ID or
all currently updating catalogs
- **Sets KB status to `ABORTED`** allowing immediate re-upload after
cancellation

**Standardizes Enum Values Across the Stack**
- **Replaces string-based status values** (`"updating"`, `"completed"`)
with strongly-typed protobuf enum constants
(`KNOWLEDGE_BASE_UPDATE_STATUS_UPDATING`,
`KNOWLEDGE_BASE_UPDATE_STATUS_COMPLETED`)
- **Adds database migration (000034)** to update all existing records to
use standardized enum values
- **Ensures consistency** between protobuf definitions, Go backend code
(via `.String()`), database storage, and integration test assertions
- **Introduces 10 well-documented lifecycle states**: UNSPECIFIED, NONE,
UPDATING, SYNCING, VALIDATING, SWAPPING, COMPLETED, FAILED, ROLLED_BACK,
ABORTED

**Fixes Synchronization Logic**
- **Excludes unstarted files** from synchronization wait - only actively
processing files (`PROCESSING`, `CONVERTING`, `CHUNKING`, `EMBEDDING`)
now block the swap
- **Prevents indefinite hangs** when files are uploaded but not yet
processed during the update lifecycle
- **Enables proper KB locking** - after lock, no new processing starts,
so unstarted files can be safely ignored

**Adds Comprehensive Test Coverage**
- **Group 5: 10 corner case tests** covering file operations during
different update phases (adding/deleting during swap, race conditions,
rapid operations, late-phase uploads)
- **Updated all test assertions** to handle both legacy string format
and new enum format for backward compatibility
- **Fixed `pollUpdateCompletion` helper** to recognize new enum-based
status values

## Architecture Highlights

### Zero-Downtime Update Strategy
1. **Phase 1-2 (UPDATING)**: Create staging KB, reprocess all files with
new config
2. **Phase 3 (SYNCING)**: Lock KB, wait for active processing to
complete (excluding unstarted files)
3. **Phase 4 (VALIDATING)**: Verify data integrity (file counts,
embeddings, chunks)
4. **Phase 5 (SWAPPING)**: Atomic pointer swap - production → rollback,
staging → production
5. **Phase 6 (COMPLETED)**: Cleanup, retain rollback KB for configured
period

### Dual-Processing During Updates
- Files uploaded during update are **processed to both production and
staging KBs**
- After swap, these files exist in correct collections without
reprocessing
- Files deleted during update are **removed from both KBs** (dual
deletion)

### Graceful Abort
- Cancels ongoing Temporal workflows via
`workflowClient.CancelWorkflow()`
- Cleans up staging KB resources (files, Milvus collection, metadata)
- Production KB remains unchanged and immediately available for new
updates
- No orphaned resources or stuck workflows

Author: pinglin

Makefile

@@ -92,7 +92,7 @@ unit-test:       				## Run unit test
 integration-test:		## Run integration tests in parallel using GNU parallel
 	@echo "Running integration tests in parallel..."
 	@parallel --halt now,fail=1 --tag --line-buffer \
-		"TEST_FOLDER_ABS_PATH=${PWD} k6 run \
+		"TEST_FOLDER_ABS_PATH=${PWD} k6 run --address=\"\" \
 		-e API_GATEWAY_PROTOCOL=${API_GATEWAY_PROTOCOL} -e API_GATEWAY_URL=${API_GATEWAY_URL} \
 		-e DB_HOST=${DB_HOST} \
 		{} --no-usage-report" ::: \

cmd/worker/main.go

@@ -168,6 +168,7 @@ func main() {
 
 	// ===== CleanupKnowledgeBaseWorkflow Activities =====
 	// Activities for cleaning up entire knowledge base resources
+	w.RegisterActivity(cw.GetInProgressFileCountActivity)       // Check for in-progress files before cleanup
 	w.RegisterActivity(cw.DeleteKBFilesFromMinIOActivity)       // Delete all KB files from MinIO
 	w.RegisterActivity(cw.DropVectorDBCollectionActivity)       // Drop Milvus collection for KB
 	w.RegisterActivity(cw.DeleteKBFileRecordsActivity)          // Delete all file records from database

config/config.yaml

@@ -26,7 +26,7 @@ database:
   timezone: Etc/UTC
   pool:
     idleconnections: 10
-    maxconnections: 30
+    maxconnections: 100
     connlifetime: 30m # In minutes, e.g., '60m'
 influxdb:
   url: http://influxdb:8086

go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/google/go-cmp v0.7.0
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1
-	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251020215452-9cefa8fe58b7
+	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251021120500-309474e9c8c4
 	github.com/instill-ai/usage-client v0.4.0
 	github.com/instill-ai/x v0.10.1-alpha
 	github.com/knadh/koanf v1.5.0

go.sum

@@ -301,8 +301,8 @@ github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpO
 github.com/hydrogen18/memlistener v0.0.0-20200120041712-dcc25e7acd91/go.mod h1:qEIFzExnS6016fRpRfxrExeVn2gbClQA99gQhnIcdhE=
 github.com/imkira/go-interpol v1.1.0/go.mod h1:z0h2/2T3XF8kyEPpRgJ3kmNv+C43p+I/CoI+jC3w2iA=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251020215452-9cefa8fe58b7 h1:fPIHRjajavlqcTdOmjUEBkar8QKxQBS4LciLJnTv6as=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251020215452-9cefa8fe58b7/go.mod h1:bCnBosofpaUxKBuTTJM3/I3thAK37kvfBnKByjnLsl4=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251021120500-309474e9c8c4 h1:tEPHqGxFwKeHEm3OhZ+WLRRmQ9BY0i1u48VNO59N0oY=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20251021120500-309474e9c8c4/go.mod h1:bCnBosofpaUxKBuTTJM3/I3thAK37kvfBnKByjnLsl4=
 github.com/instill-ai/usage-client v0.4.0 h1:xf1hAlO4a8lZwZzz9bprZOJqU3ghIcIsavUUB7UURyg=
 github.com/instill-ai/usage-client v0.4.0/go.mod h1:zZ9LRoXps2u63ARYPAbR2YvqTib3dWJLObZn+9YqhF0=
 github.com/instill-ai/x v0.10.1-alpha h1:15+SntJqoh1ab+QhS8OdoagyUMhR8NU4UV12Xt43qmo=

integration-test/const.js

@@ -17,13 +17,15 @@ if (__ENV.API_GATEWAY_PROTOCOL) {
   proto = "http"
 }
 
-export const artifactPrivateHost = `http://artifact-backend:3082`;
-export const artifactPublicHost = apiGatewayMode ? `${proto}://${__ENV.API_GATEWAY_URL}` : `http://api-gateway:8080`
-export const mgmtPublicHost = apiGatewayMode ? `${proto}://${__ENV.API_GATEWAY_URL}` : `http://api-gateway:8080`
-export const artifactGRPCPrivateHost = `artifact-backend:3082`;
+export const artifactRESTPublicHost = apiGatewayMode ? `${proto}://${__ENV.API_GATEWAY_URL}` : `http://api-gateway:8080`
+export const mgmtRESTPublicHost = apiGatewayMode ? `${proto}://${__ENV.API_GATEWAY_URL}` : `http://api-gateway:8080`
+
 export const artifactGRPCPublicHost = apiGatewayMode ? `${__ENV.API_GATEWAY_URL}` : `api-gateway:8080`;
+export const artifactGRPCPrivateHost = apiGatewayMode ? `localhost:3082` : `artifact-backend:3082`;
+
 export const mgmtGRPCPublicHost = apiGatewayMode ? `${__ENV.API_GATEWAY_URL}` : `api-gateway:8080`;
-export const mgmtGRPCPrivateHost = `mgmt-backend:3084`;
+export const mgmtGRPCPrivateHost = apiGatewayMode ? `localhost:3084` : `mgmt-backend:3084`;
+
 export const mgmtVersion = 'v1beta';
 
 export const namespace = "users/admin"

integration-test/grpc-kb-update.js

@@ -8,16 +8,16 @@ export function CheckListRepositoryTags(client, data) {
     check(true, { [constant.banner(groupName)]: () => true });
 
     var resOrigin = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTags",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTagsAdmin",
       {
         parent: `repositories/${data.expectedOwner.id}`,
         pageSize: 10
       },
       data.metadata
     );
     check(resOrigin, {
-      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTags response status is StatusOK": (r) => r.status === grpc.StatusOK,
-      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTags response tags is array": (r) =>
+      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTagsAdmin response status is StatusOK": (r) => r.status === grpc.StatusOK,
+      "artifact.artifact.v1alpha.ArtifactPrivateService/ListRepositoryTagsAdmin response tags is array": (r) =>
         Array.isArray(r.message.tags),
     });
   });
@@ -35,7 +35,7 @@ export function CheckGetRepositoryTag(client, data) {
     }
 
     var res = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/GetRepositoryTag",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/GetRepositoryTagAdmin",
       {
         name: tagName,
       },
@@ -58,7 +58,7 @@ export function CheckCreateAndDeleteRepositoryTag(client, data) {
     const tagName = `repositories/${repoId}/tags/${tagId}`;
 
     var createRes = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/CreateRepositoryTag",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/CreateRepositoryTagAdmin",
       {
         tag: {
           name: tagName,
@@ -72,7 +72,7 @@ export function CheckCreateAndDeleteRepositoryTag(client, data) {
     });
 
     var deleteRes = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/DeleteRepositoryTag",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/DeleteRepositoryTagAdmin",
       {
         name: tagName,
       },
@@ -94,7 +94,7 @@ export function CheckGetObject(client, data) {
       return;
     }
     var res = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/GetObject",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/GetObjectAdmin",
       { uid: uid },
       data.metadata
     );
@@ -116,7 +116,7 @@ export function CheckGetObjectURL(client, data) {
       return;
     }
     var res = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/GetObjectURL",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/GetObjectURLAdmin",
       uid ? { uid: uid } : { encodedUrlPath: encoded },
       data.metadata
     );
@@ -136,7 +136,7 @@ export function CheckUpdateObject(client, data) {
       return;
     }
     var res = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/UpdateObject",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/UpdateObjectAdmin",
       { uid: uid, isUploaded: true },
       data.metadata
     );
@@ -156,7 +156,7 @@ export function CheckGetFileAsMarkdown(client, data) {
       return;
     }
     var res = client.invoke(
-      "artifact.artifact.v1alpha.ArtifactPrivateService/GetFileAsMarkdown",
+      "artifact.artifact.v1alpha.ArtifactPrivateService/GetFileAsMarkdownAdmin",
       { fileUid: fileUid },
       data.metadata
     );

integration-test/grpc-private.js

@@ -218,8 +218,8 @@ export function CheckCleanupOnFileDeletion(client, data) {
     const checkAfter = `
       SELECT
         (SELECT COUNT(*) FROM converted_file WHERE file_uid = '${file.fileUid}') as converted,
-        (SELECT COUNT(*) FROM text_chunk WHERE kb_file_uid = '${file.fileUid}') as chunks,
-        (SELECT COUNT(*) FROM embedding WHERE kb_file_uid = '${file.fileUid}') as embeddings
+        (SELECT COUNT(*) FROM text_chunk WHERE file_uid = '${file.fileUid}') as chunks,
+        (SELECT COUNT(*) FROM embedding WHERE file_uid = '${file.fileUid}') as embeddings
     `;
     try {
       const result = constant.db.exec(checkAfter);

integration-test/grpc-public.js

@@ -3,7 +3,7 @@ import http from "k6/http";
 
 import { check, group } from "k6";
 
-import * as catalogPrivate from "./grpc-private.js";
+import * as grpcPrivate from "./grpc-private.js";
 import * as grpcPublic from "./grpc-public.js";
 import * as grpcPublicWithJwt from "./grpc-public-with-jwt.js";
 
@@ -39,22 +39,22 @@ export function setup() {
 
   // Clean up any leftover test data from previous runs
   try {
-    constant.db.exec(`DELETE FROM text_chunk WHERE kb_file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
-    constant.db.exec(`DELETE FROM embedding WHERE kb_file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
+    constant.db.exec(`DELETE FROM text_chunk WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
+    constant.db.exec(`DELETE FROM embedding WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
     constant.db.exec(`DELETE FROM converted_file WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
     constant.db.exec(`DELETE FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%'`);
     constant.db.exec(`DELETE FROM knowledge_base WHERE id LIKE '${constant.dbIDPrefix}%'`);
   } catch (e) {
     console.log(`Setup cleanup warning: ${e}`);
   }
 
-  var loginResp = http.request("POST", `${constant.mgmtPublicHost}/v1beta/auth/login`, JSON.stringify({
+  var loginResp = http.request("POST", `${constant.mgmtRESTPublicHost}/v1beta/auth/login`, JSON.stringify({
     "username": constant.defaultUsername,
     "password": constant.defaultPassword,
   }))
 
   check(loginResp, {
-    [`POST ${constant.mgmtPublicHost}/v1beta/auth/login response status is 200`]: (
+    [`POST ${constant.mgmtRESTPublicHost}/v1beta/auth/login response status is 200`]: (
       r
     ) => r.status === 200,
   });
@@ -115,13 +115,13 @@ export default function (data) {
     privateClient.connect(constant.artifactGRPCPrivateHost, {
       plaintext: true,
     });
-    catalogPrivate.CheckListRepositoryTags(privateClient, data);
-    catalogPrivate.CheckGetRepositoryTag(privateClient, data);
-    catalogPrivate.CheckCreateAndDeleteRepositoryTag(privateClient, data);
-    catalogPrivate.CheckGetObject(privateClient, data);
-    catalogPrivate.CheckGetObjectURL(privateClient, data);
-    catalogPrivate.CheckUpdateObject(privateClient, data);
-    catalogPrivate.CheckGetFileAsMarkdown(privateClient, data);
+    grpcPrivate.CheckListRepositoryTags(privateClient, data);
+    grpcPrivate.CheckGetRepositoryTag(privateClient, data);
+    grpcPrivate.CheckCreateAndDeleteRepositoryTag(privateClient, data);
+    grpcPrivate.CheckGetObject(privateClient, data);
+    grpcPrivate.CheckGetObjectURL(privateClient, data);
+    grpcPrivate.CheckUpdateObject(privateClient, data);
+    grpcPrivate.CheckGetFileAsMarkdown(privateClient, data);
     privateClient.close();
   }
 
@@ -156,8 +156,8 @@ export function teardown(data) {
     check(true, { [constant.banner(groupName)]: () => true });
 
     // Delete from child tables first, before deleting parent records
-    constant.db.exec(`DELETE FROM text_chunk WHERE kb_file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
-    constant.db.exec(`DELETE FROM embedding WHERE kb_file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
+    constant.db.exec(`DELETE FROM text_chunk WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
+    constant.db.exec(`DELETE FROM embedding WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
     constant.db.exec(`DELETE FROM converted_file WHERE file_uid IN (SELECT uid FROM knowledge_base_file WHERE name LIKE '${constant.dbIDPrefix}%')`);
 
     // Now delete parent tables