fix(gateway): route artifact task lookup through task access

vincentkoc · vincentkoc · commit 188a5005c7a1 · 2026-04-30T00:58:28.000-07:00
diff --git a/src/gateway/server-methods/artifacts.test.ts b/src/gateway/server-methods/artifacts.test.ts
@@ -2,14 +2,14 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { artifactsHandlers, collectArtifactsFromMessages } from "./artifacts.js";
 
 const hoisted = vi.hoisted(() => ({
-  getTaskById: vi.fn(),
+  getTaskSessionLookupByIdForStatus: vi.fn(),
   loadSessionEntry: vi.fn(),
   readSessionMessages: vi.fn(),
   resolveSessionKeyForRun: vi.fn(),
 }));
 
-vi.mock("../../tasks/task-registry.js", () => ({
-  getTaskById: hoisted.getTaskById,
+vi.mock("../../tasks/task-status-access.js", () => ({
+  getTaskSessionLookupByIdForStatus: hoisted.getTaskSessionLookupByIdForStatus,
 }));
 
 vi.mock("../session-utils.js", async () => {
@@ -44,7 +44,7 @@ function createResponder() {
 describe("artifacts RPC handlers", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    hoisted.getTaskById.mockReturnValue(undefined);
+    hoisted.getTaskSessionLookupByIdForStatus.mockReturnValue(undefined);
     hoisted.loadSessionEntry.mockReturnValue({
       storePath: "/tmp/sessions.json",
       entry: { sessionId: "sess-main", sessionFile: "/tmp/sess-main.jsonl" },
@@ -161,9 +161,8 @@ describe("artifacts RPC handlers", () => {
     expect(payload.artifacts?.[0]).toMatchObject({ runId: "run-1" });
   });
 
-  it("resolves taskId queries through the task registry and filters artifacts by messageTaskId", async () => {
-    hoisted.getTaskById.mockReturnValue({
-      taskId: "task-1",
+  it("resolves taskId queries through task status access and filters artifacts by messageTaskId", async () => {
+    hoisted.getTaskSessionLookupByIdForStatus.mockReturnValue({
       requesterSessionKey: "agent:main:main",
       runId: "run-for-task-1",
     });
@@ -196,7 +195,7 @@ describe("artifacts RPC handlers", () => {
     });
 
     expect(list.calls[0]?.ok).toBe(true);
-    expect(hoisted.getTaskById).toHaveBeenCalledWith("task-1");
+    expect(hoisted.getTaskSessionLookupByIdForStatus).toHaveBeenCalledWith("task-1");
     expect(hoisted.resolveSessionKeyForRun).not.toHaveBeenCalled();
     expect(hoisted.loadSessionEntry).toHaveBeenCalledWith("agent:main:main");
     const listPayload = list.calls[0]?.payload as { artifacts?: Array<Record<string, unknown>> };
diff --git a/src/gateway/server-methods/artifacts.ts b/src/gateway/server-methods/artifacts.ts
@@ -1,5 +1,5 @@
 import { createHash } from "node:crypto";
-import { getTaskById } from "../../tasks/task-registry.js";
+import { getTaskSessionLookupByIdForStatus } from "../../tasks/task-status-access.js";
 import {
   ErrorCodes,
   errorShape,
@@ -278,7 +278,7 @@ function resolveQuerySessionKey(query: ArtifactQuery): string | undefined {
     return resolveSessionKeyForRun(query.runId);
   }
   if (query.taskId) {
-    const task = getTaskById(query.taskId);
+    const task = getTaskSessionLookupByIdForStatus(query.taskId);
     const requesterSessionKey = asNonEmptyString(task?.requesterSessionKey);
     if (requesterSessionKey) {
       return requesterSessionKey;
diff --git a/src/tasks/task-status-access.ts b/src/tasks/task-status-access.ts
@@ -1,6 +1,18 @@
-import { listTasksForAgentId, listTasksForSessionKey } from "./task-registry.js";
+import { getTaskById, listTasksForAgentId, listTasksForSessionKey } from "./task-registry.js";
 import type { TaskRecord } from "./task-registry.types.js";
 
+export function getTaskSessionLookupByIdForStatus(
+  taskId: string,
+): Pick<TaskRecord, "requesterSessionKey" | "runId"> | undefined {
+  const task = getTaskById(taskId);
+  return task
+    ? {
+        requesterSessionKey: task.requesterSessionKey,
+        ...(task.runId ? { runId: task.runId } : {}),
+      }
+    : undefined;
+}
+
 export function listTasksForSessionKeyForStatus(sessionKey: string): TaskRecord[] {
   return listTasksForSessionKey(sessionKey);
 }
