Fix data task runtime restore and datasource readiness

Author: Wistomize

apps/web/src/app/data-tasks/__tests__/conversation-restore.test.ts

@@ -18,6 +18,7 @@ import {
 import { ConfigApiError } from "../../../lib/config-api/types";
 import {
   createInitialLiveRun,
+  deriveRunUsage,
   formatWorkspaceMetadataSummary,
   reconcileLiveRunArtifacts,
   reduceLiveRunEvent,
@@ -1331,6 +1332,93 @@ describe("hydrateLiveRunFromConversation run ordering", () => {
 });
 
 describe("hydrateLiveRunFromConversation", () => {
+  it("restores canceled message-only checkpoints as canceled live runs", () => {
+    const dto: SessionConversationDto = {
+      sessionId: "thread-1",
+      messages: [
+        {
+          id: "run-canceled:user",
+          runId: "run-canceled",
+          role: "user",
+          source: "client",
+          messageId: "frontend-user-canceled",
+          contentText: "Please write a long answer.",
+          position: 1,
+          createdAt: "2026-06-25T10:00:01Z",
+        },
+        {
+          id: "run-canceled:assistant:partial",
+          runId: "run-canceled",
+          role: "assistant",
+          source: "agent",
+          messageId: "assistant-partial",
+          contentText: "Partial answer before cancellation.",
+          position: 2,
+          createdAt: "2026-06-25T10:00:02Z",
+        },
+      ],
+      runEventRefs: [{ runId: "run-canceled", eventCount: 12, firstSeq: 1, lastSeq: 12 }],
+      checkpoints: [
+        {
+          runId: "run-canceled",
+          status: "canceled",
+          messageStartPosition: 1,
+          messageEndPosition: 2,
+          firstEventSeq: 1,
+          lastEventSeq: 12,
+          startedAt: "2026-06-25T10:00:00Z",
+          finishedAt: "2026-06-25T10:00:03Z",
+          errorMessage: "user-requested",
+        },
+      ],
+      toolCalls: [],
+    };
+
+    const run = hydrateLiveRunFromConversation(createInitialLiveRun(), dto);
+    expect(run.runStatus).toBe("canceled");
+    expect(run.runStartedAt).toBe(Date.parse("2026-06-25T10:00:00Z"));
+    expect(run.runFinishedAt).toBe(Date.parse("2026-06-25T10:00:03Z"));
+    expect(deriveRunUsage(run).durationMs).toBe(3000);
+  });
+
+  it("restores canceled tool checkpoints with checkpoint timing", () => {
+    const dto: SessionConversationDto = {
+      sessionId: "thread-1",
+      messages: [],
+      runEventRefs: [{ runId: "run-canceled", eventCount: 4, firstSeq: 1, lastSeq: 4 }],
+      checkpoints: [
+        {
+          runId: "run-canceled",
+          status: "canceled",
+          messageStartPosition: 1,
+          messageEndPosition: 2,
+          firstEventSeq: 1,
+          lastEventSeq: 4,
+          startedAt: "2026-06-25T10:00:00Z",
+          finishedAt: "2026-06-25T10:00:05Z",
+          errorMessage: "user-requested",
+        },
+      ],
+      toolCalls: [
+        {
+          runId: "run-canceled",
+          toolCallId: "sql-canceled",
+          status: "completed",
+          toolName: "run_sql_readonly",
+          callEventSeq: 2,
+          resultEventSeq: 3,
+          resultPreview: JSON.stringify({ row_count: 3 }),
+        },
+      ],
+    };
+
+    const run = hydrateLiveRunFromConversation(createInitialLiveRun(), dto);
+    expect(run.runStatus).toBe("canceled");
+    expect(run.runStartedAt).toBe(Date.parse("2026-06-25T10:00:00Z"));
+    expect(run.runFinishedAt).toBe(Date.parse("2026-06-25T10:00:05Z"));
+    expect(deriveRunUsage(run).durationMs).toBe(5000);
+  });
+
   it("hydrates live run tool calls and links restored artifacts", () => {
     const dto: SessionConversationDto = {
       sessionId: "thread-1",

apps/web/src/app/data-tasks/__tests__/per-run-mentions.test.ts

@@ -20,8 +20,12 @@ import {
   type WorkspaceConfigStore,
 } from "../data-task-state";
 
-function item(id: string, name = id) {
-  return { id, name, description: `${name} desc`, enabled: true };
+function item(
+  id: string,
+  name = id,
+  status: "connected" | "failed" | "untested" = "connected",
+) {
+  return { id, name, description: `${name} desc`, enabled: true, status };
 }
 
 const store: WorkspaceConfigStore = {
@@ -179,6 +183,39 @@ describe("buildRunConfig", () => {
     expect(config.activeDatasourceId).toBe("db-orders");
   });
 
+  it("excludes unusable datasources from run config and active fallback", () => {
+    const mixedStore: WorkspaceConfigStore = {
+      ...store,
+      db: [
+        item("db-failed", "Failed db", "failed"),
+        item("db-default"),
+        item("db-untested", "Untested db", "untested"),
+      ],
+    };
+    const selection = togglePerRunMention(
+      emptyPerRunSelection(),
+      "db",
+      "db-failed",
+    );
+    const config = buildRunConfig(mixedStore, {
+      activeLlmId: "llm-1",
+      defaultDatasourceId: "db-failed",
+      session,
+      perRunSelection: selection,
+    });
+    expect(config.enabledDatasourceIds).toEqual(["db-default"]);
+    expect(config.activeDatasourceId).toBe("db-default");
+    expect(config.mentioned.db).toEqual([]);
+    expect(
+      resolveActiveDatasourceId(
+        mixedStore,
+        session,
+        selection,
+        "db-failed",
+      ),
+    ).toBe("db-default");
+  });
+
   it("splits file mentions into workspace fileIds and session pinnedPaths", () => {
     const files: FileMentionResource[] = [
       {

apps/web/src/app/data-tasks/__tests__/session-pane-ui.test.ts

@@ -11,8 +11,8 @@ import {
 import type { WorkspaceConfigStore } from "../data-task-state";
 import { LEFT_PANEL_MAX_WIDTH } from "../workspace-layout";
 
-function item(id: string, name = id) {
-  return { id, name, description: `${name} desc`, enabled: true };
+function item(id: string, name = id, status: "connected" | "failed" | "untested" = "connected") {
+  return { id, name, description: `${name} desc`, enabled: true, status };
 }
 
 const workspaceConfig: WorkspaceConfigStore = {
@@ -127,6 +127,32 @@ describe("session pane ui conventions", () => {
     });
   });
 
+  it("counts only runnable data sources in the workspace rail", () => {
+    const groups = getWorkspaceResourceNavGroups({
+      workspaceConfig: {
+        ...workspaceConfig,
+        db: [
+          item("sqlite-orders", "Orders SQLite", "connected"),
+          item("api-duckdb-demo", "API DuckDB Demo", "failed"),
+          item("warehouse", "Warehouse", "untested"),
+        ],
+      },
+      workspaceFileCount: 0,
+      activeConfigPanel: null,
+      activeFilesPanel: false,
+      capabilitiesReady: true,
+      supportsFiles: true,
+      supportsKnowledge: true,
+      supportsMcp: true,
+      supportsSkills: true,
+    });
+
+    expect(groups[0]).toMatchObject({
+      title: "Data Sources",
+      summary: "1",
+    });
+  });
+
   it("keeps sidebar status minimal for static configuration navigation", () => {
     const groups = getWorkspaceResourceNavGroups({
       workspaceConfig: { ...workspaceConfig, db: [] },

apps/web/src/app/data-tasks/conversation-restore.ts

@@ -613,6 +613,14 @@ function checkpointByRunId(
   return map;
 }
 
+function timestampFromIso(value: string | undefined): number | undefined {
+  if (!value) {
+    return undefined;
+  }
+  const parsed = Date.parse(value);
+  return Number.isFinite(parsed) ? parsed : undefined;
+}
+
 function runStatusFromCheckpoint(
   checkpoint: ConversationCheckpointDto | undefined,
 ): LiveRun["runStatus"] | undefined {
@@ -1273,11 +1281,25 @@ function applyConversationToolCall(state: LiveRun, toolCall: ConversationToolCal
 function finalizeHydratedRunSegment(
   state: LiveRun,
   tools: ConversationToolCallDto[],
+  checkpoint: ConversationCheckpointDto | undefined,
 ): LiveRun {
+  const checkpointStatus = runStatusFromCheckpoint(checkpoint);
+  if (checkpointStatus === "running" || checkpointStatus === "queued") {
+    return state;
+  }
+  if (checkpointStatus === "failed") {
+    return reduceLiveRunEvent(state, {
+      type: "RUN_ERROR",
+      message: checkpoint?.errorMessage?.trim() || ORPHANED_RUN_ASSISTANT_PLACEHOLDER,
+    });
+  }
+  if (checkpointStatus === "canceled") {
+    return reduceLiveRunEvent(state, { type: "RUN_FINISHED", status: "canceled" });
+  }
   if (tools.length === 0) {
     return state;
   }
-  if (isPendingCollaborationRunEnd(tools)) {
+  if (checkpointStatus === "suspended" || isPendingCollaborationRunEnd(tools)) {
     let next = reduceLiveRunEvent(state, {
       type: "STATE_DELTA",
       delta: [{ op: "replace", path: "/runStatus", value: "suspended" }],
@@ -1291,9 +1313,9 @@ function finalizeHydratedRunSegment(
 function finalizeMessageOnlyHydratedRunSegment(
   state: LiveRun,
   segment: RestorableRunSegment,
+  checkpoint: ConversationCheckpointDto | undefined,
   dto: SessionConversationDto,
 ): LiveRun {
-  const checkpoint = checkpointByRunId(dto).get(segment.runId);
   const checkpointStatus = runStatusFromCheckpoint(checkpoint);
   if (checkpointStatus === "running" || checkpointStatus === "queued") {
     return state;
@@ -1305,7 +1327,7 @@ function finalizeMessageOnlyHydratedRunSegment(
     });
   }
   if (checkpointStatus === "canceled") {
-    return reduceLiveRunEvent(state, { type: "RUN_FINISHED" });
+    return reduceLiveRunEvent(state, { type: "RUN_FINISHED", status: "canceled" });
   }
   if (segment.hasUser && !segment.hasAssistant) {
     return reduceLiveRunEvent(state, {
@@ -1316,6 +1338,22 @@ function finalizeMessageOnlyHydratedRunSegment(
   return reduceLiveRunEvent(state, { type: "RUN_FINISHED" });
 }
 
+function applyHydratedRunCheckpointTiming(
+  state: LiveRun,
+  checkpoint: ConversationCheckpointDto | undefined,
+): LiveRun {
+  const startedAt = timestampFromIso(checkpoint?.startedAt);
+  const finishedAt = timestampFromIso(checkpoint?.finishedAt);
+  if (startedAt === undefined && finishedAt === undefined) {
+    return state;
+  }
+  return {
+    ...state,
+    ...(startedAt !== undefined ? { runStartedAt: startedAt } : {}),
+    ...(finishedAt !== undefined ? { runFinishedAt: finishedAt } : {}),
+  };
+}
+
 function startNextHydratedRunGroup(state: LiveRun, runId?: string): LiveRun {
   if (state.runStatus === "suspended") {
     const archived: LiveRun = {
@@ -1470,9 +1508,11 @@ export function hydrateLiveRunFromConversation(
   }
 
   const runSegments = collectRestorableRunSegments(dto);
+  const checkpoints = checkpointByRunId(dto);
 
   let next = createInitialLiveRun();
   for (const [index, segment] of runSegments.entries()) {
+    const checkpoint = checkpoints.get(segment.runId);
     next =
       index === 0
         ? reduceLiveRunEvent(next, { type: "RUN_STARTED", runId: segment.runId })
@@ -1486,8 +1526,9 @@ export function hydrateLiveRunFromConversation(
 
     next =
       segment.tools.length > 0
-        ? finalizeHydratedRunSegment(next, segment.tools)
-        : finalizeMessageOnlyHydratedRunSegment(next, segment, dto);
+        ? finalizeHydratedRunSegment(next, segment.tools, checkpoint)
+        : finalizeMessageOnlyHydratedRunSegment(next, segment, checkpoint, dto);
+    next = applyHydratedRunCheckpointTiming(next, checkpoint);
   }
 
   next = reconcileLiveRunArtifacts(mergePreservedLiveRunSessionData(state, next));