feat(runtime-host): write Read images to Session context

Author: likun

packages/runtime-host/src/__tests__/execution-composition.test.ts

@@ -98,7 +98,7 @@ test('production composition owns the long-term memory database lifecycle', asyn
   });
 });
 
-test('production composition reaches Ready when the optional context reader cannot open', async () => {
+test('production composition reaches Ready when the optional context Store cannot open', async () => {
   await withCompositionRoot(async ({ root, owner }) => {
     await mkdir(join(root, CONTEXT_OFFLOAD_DATABASE_NAME));
     const originalConsoleError = console.error;
@@ -109,7 +109,7 @@ test('production composition reaches Ready when the optional context reader cann
       composition = await createExecutionRuntimeHostComposition(compositionContext(owner));
       assert.equal(composition.workspaceExecution.state, 'ready');
       assert.equal(
-        diagnostics.some((message) => message.includes('optional context-offload reader')),
+        diagnostics.some((message) => message.includes('optional context-offload Store')),
         true,
       );
     } finally {

packages/runtime-host/src/__tests__/protocol.test.ts

@@ -134,6 +134,10 @@ describe('Runtime Host bootstrap protocol', () => {
     assert.ok(RUNTIME_HOST_COMPATIBILITY_EPOCH > 22);
   });
 
+  test('publishes a new compatibility epoch for Read image Session context refs', () => {
+    assert.ok(RUNTIME_HOST_COMPATIBILITY_EPOCH > 67);
+  });
+
   test('rejects the legacy connection update result in the current compatibility epoch', () => {
     assert.throws(
       () =>

packages/runtime-host/src/__tests__/session-retirement-coordinator.test.ts

@@ -228,7 +228,7 @@ describe('Host Session retirement coordinator', () => {
         'parent retirement cleanup did not converge',
       );
       assert.deepEqual(new Set(harness.actions.purgedArtifacts), new Set(harness.familyIds));
-      assert.deepEqual(new Set(harness.actions.checkedContext), new Set(harness.familyIds));
+      assert.deepEqual(new Set(harness.actions.retiredContext), new Set(harness.familyIds));
     });
   });
 
@@ -1009,7 +1009,7 @@ interface RetirementActions {
   readonly retiredCapabilities: string[];
   readonly retiredMessages: string[];
   readonly purgedArtifacts: string[];
-  readonly checkedContext: string[];
+  readonly retiredContext: string[];
   readonly purgedTasks: string[];
   readonly purgedOperationalState: string[];
   readonly purgedAgentGraphs: string[];
@@ -1048,7 +1048,7 @@ async function withHarness(
       retiredCapabilities: [],
       retiredMessages: [],
       purgedArtifacts: [],
-      checkedContext: [],
+      retiredContext: [],
       purgedTasks: [],
       purgedOperationalState: [],
       purgedAgentGraphs: [],
@@ -1214,8 +1214,11 @@ async function withHarness(
           actions.purgedTasks.push(sessionId);
         },
       },
-      assertNoContextOffloadReferences: async (sessionIds) => {
-        actions.checkedContext.push(...sessionIds);
+      contextOffload: {
+        retireSession: async (sessionId) => {
+          actions.retiredContext.push(sessionId);
+          return { releasedReferences: 0, releasedLogicalBytes: 0 };
+        },
       },
       purgeOperationalState: async (sessionId) => {
         actions.purgedOperationalState.push(sessionId);

packages/runtime-host/src/protocol/index.ts

@@ -94,7 +94,8 @@ export const RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION = 1 as const;
 export const RUNTIME_HOST_PROTOCOL_VERSION = 0 as const;
 // Increment when the same protocol version no longer guarantees safe Client-Host
 // interoperability. Mismatches are rejected before domain commands are admitted.
-export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 67 as const;
+export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 68 as const;
+// 68: Read image tool results may carry durable `session_context` refs.
 // 67: Message lifecycle queries expose durable execution ownership and
 // cancellation. Older peers cannot decode or provide the closed proof list.
 // 66: Peer Mesh queries expose one canonical transit selection and runtime metrics.

packages/runtime-host/src/server/execution-composition.ts

@@ -77,14 +77,9 @@ import {
 import { type MakaTool } from '@maka/runtime/tool-runtime';
 import { type RuntimeHostedRootAuthority } from '@maka/runtime/message-authority';
 import { createAgentGraphControlStore } from '@maka/storage/agent-graph-control-store';
-import {
-  createArtifactAttachmentResourceReader,
-  createReadImageSnapshotter,
-} from '@maka/storage/artifact-stores';
-import {
-  isSessionNotFoundError,
-  SessionMetadataConflictError,
-} from '@maka/storage/execution-stores';
+import { createArtifactAttachmentResourceReader } from '@maka/storage/artifact-stores';
+import { createReadImageSnapshotStore } from '@maka/storage/read-image-snapshot-store';
+import { isSessionNotFoundError } from '@maka/storage/execution-stores';
 import { createExternalSessionAdapterRegistry } from '@maka/storage/external-sessions';
 import { createGitWorktreeChildExecutor } from '@maka/storage/git-worktree-child-executor';
 import { runWithStorageRootLease } from '@maka/storage/root-authority';
@@ -202,15 +197,16 @@ export interface ExecutionRuntimeHostComposition extends RuntimeHostComposition
   readonly workspaceExecution: RuntimeHostWorkspaceExecutionComposition;
 }
 
-const CONTEXT_OFFLOAD_READER_LIMITS: ContextOffloadLimits = Object.freeze({
+const GIBIBYTE = 1024 * 1024 * 1024;
+const CONTEXT_OFFLOAD_LIMITS: ContextOffloadLimits = Object.freeze({
   ownerMaxBytes: Object.freeze({
     read_image_snapshot: MAX_READ_IMAGE_BYTES,
     tool_result_archive: 0,
   }),
-  // This expand slice opens only the reader path. Zero quotas make accidental
-  // non-empty puts fail closed until the writer/lifecycle cutover lands.
-  sessionLogicalBytes: 0,
-  workspacePhysicalBytes: 0,
+  // Read images are bounded individually and logically per Session. Physical
+  // bytes are content-addressed across Sessions and bounded per workspace.
+  sessionLogicalBytes: GIBIBYTE,
+  workspacePhysicalBytes: 20 * GIBIBYTE,
 });
 
 export interface CreateExecutionRuntimeHostCompositionOptions {
@@ -239,7 +235,7 @@ export async function createExecutionRuntimeHostComposition(
   dependencies: ExecutionRuntimeHostCompositionDependencies = {},
 ): Promise<ExecutionRuntimeHostComposition> {
   const storage = await openStorageWriterComposition(context.owner.lease, {
-    contextOffloadLimits: CONTEXT_OFFLOAD_READER_LIMITS,
+    contextOffloadLimits: CONTEXT_OFFLOAD_LIMITS,
     afterRuntimePolicyOpened: async (stores) => {
       if (options.bootstrapRuntimePolicy !== false) {
         await ensureBootstrapRuntimePolicy({
@@ -255,7 +251,7 @@ export async function createExecutionRuntimeHostComposition(
   });
   if (storage.contextOffloadUnavailable) {
     console.error(
-      `[runtime-host] optional context-offload reader could not be opened: ${generalizedErrorMessage(storage.contextOffloadUnavailable.cause)}`,
+      `[runtime-host] optional context-offload Store could not be opened: ${generalizedErrorMessage(storage.contextOffloadUnavailable.cause)}`,
     );
   }
   const stores = storage.execution;
@@ -285,6 +281,17 @@ export async function createExecutionRuntimeHostComposition(
     const openedContextOffloadReader = openedContextOffloadStore
       ? createInteractiveContextOffloadReader(openedContextOffloadStore)
       : undefined;
+    const contextOffloadRetirement = openedContextOffloadStore
+      ? openedContextOffloadStore
+      : storage.contextOffloadUnavailable
+        ? {
+            retireSession: async (_sessionId: string): Promise<never> => {
+              throw new Error('Context-offload Store is unavailable during Session retirement', {
+                cause: storage.contextOffloadUnavailable?.cause,
+              });
+            },
+          }
+        : undefined;
     const openedUsageStores = storage.usage;
     const openedShellRunStore = storage.shellRuns;
     const worktreeChildExecutor = createGitWorktreeChildExecutor({
@@ -395,7 +402,21 @@ export async function createExecutionRuntimeHostComposition(
       }),
       backgroundTasks: runtimeResources,
       ptyControls: runtimeResources,
-      snapshotImage: createReadImageSnapshotter(openedArtifactStore),
+      ...(openedContextOffloadStore
+        ? {
+            snapshotImage: async (input: {
+              readonly sessionId: string;
+              readonly ownerId: string;
+              readonly bytes: Uint8Array;
+              readonly mimeType: string;
+            }) =>
+              createReadImageSnapshotStore(openedContextOffloadStore, input.sessionId).snapshot({
+                ownerId: input.ownerId,
+                bytes: input.bytes,
+                mimeType: input.mimeType,
+              }),
+          }
+        : {}),
       ...(sandboxManager ? { sandboxManager } : {}),
       ...(filesystemWorker ? { filesystemWorker } : {}),
     };
@@ -1437,6 +1458,7 @@ export async function createExecutionRuntimeHostComposition(
       stores,
       artifacts: openedArtifactStore,
       taskLedger: taskLedgerStore,
+      ...(openedContextOffloadStore ? { contextOffload: openedContextOffloadStore } : {}),
       manager,
       admission: sessionAdmission,
       continuity: continuityCoordinator,
@@ -1461,20 +1483,7 @@ export async function createExecutionRuntimeHostComposition(
       continuity: continuityCoordinator,
       artifacts: openedArtifactStore,
       taskLedger: taskLedgerStore,
-      assertNoContextOffloadReferences: async (sessionIds) => {
-        if (!openedContextOffloadStore) {
-          throw new Error('Context-offload reader is unavailable during Session removal', {
-            cause: storage.contextOffloadUnavailable?.cause,
-          });
-        }
-        for (const sessionId of sessionIds) {
-          if ((await openedContextOffloadStore.usage(sessionId)).references > 0) {
-            throw new SessionMetadataConflictError(
-              'Session removal does not support Session context references yet',
-            );
-          }
-        }
-      },
+      ...(contextOffloadRetirement ? { contextOffload: contextOffloadRetirement } : {}),
       purgeOperationalState: async (sessionId) => {
         await stores.purgeConversationOperationalState(sessionId);
         await openedPlanStore.purgeSessionState(sessionId);

packages/runtime-host/src/server/session-retirement-coordinator.ts

@@ -32,6 +32,7 @@ import {
 } from '@maka/storage/execution-stores';
 import { type SessionManager } from '@maka/runtime/session-manager';
 import type { InteractiveTaskLedgerWriter } from '@maka/storage/task-ledger-authority';
+import type { InteractiveContextOffloadWriter } from '@maka/storage/context-offload-store';
 import {
   type OperationOutcome,
   type SessionCatalogItem,
@@ -121,7 +122,7 @@ export interface HostSessionRetirementCoordinatorOptions {
   readonly continuity: RetirementContinuity;
   readonly artifacts: Pick<InteractiveArtifactStoreWriter, 'purgeSessionArtifacts'>;
   readonly taskLedger: Pick<InteractiveTaskLedgerWriter, 'purgeConversationTaskLedger'>;
-  readonly assertNoContextOffloadReferences?: (sessionIds: readonly string[]) => Promise<void>;
+  readonly contextOffload?: Pick<InteractiveContextOffloadWriter, 'retireSession'>;
   readonly purgeOperationalState: (sessionId: string) => Promise<void>;
   readonly purgeAgentGraphState: (sessionId: string) => Promise<void>;
   readonly worktrees?: Pick<SubagentWorktreeExecutor, 'retire'>;
@@ -192,7 +193,7 @@ export class HostSessionRetirementCoordinator {
   readonly #continuity: RetirementContinuity;
   readonly #artifacts: HostSessionRetirementCoordinatorOptions['artifacts'];
   readonly #taskLedger: HostSessionRetirementCoordinatorOptions['taskLedger'];
-  readonly #assertNoContextOffloadReferences: HostSessionRetirementCoordinatorOptions['assertNoContextOffloadReferences'];
+  readonly #contextOffload: HostSessionRetirementCoordinatorOptions['contextOffload'];
   readonly #purgeOperationalState: HostSessionRetirementCoordinatorOptions['purgeOperationalState'];
   readonly #purgeAgentGraphState: HostSessionRetirementCoordinatorOptions['purgeAgentGraphState'];
   readonly #worktrees: HostSessionRetirementCoordinatorOptions['worktrees'];
@@ -220,7 +221,7 @@ export class HostSessionRetirementCoordinator {
     this.#continuity = options.continuity;
     this.#artifacts = options.artifacts;
     this.#taskLedger = options.taskLedger;
-    this.#assertNoContextOffloadReferences = options.assertNoContextOffloadReferences;
+    this.#contextOffload = options.contextOffload;
     this.#purgeOperationalState = options.purgeOperationalState;
     this.#purgeAgentGraphState = options.purgeAgentGraphState;
     this.#worktrees = options.worktrees;
@@ -336,7 +337,6 @@ export class HostSessionRetirementCoordinator {
           if (plan.archive.sessionIds.length > 0) {
             archiveHandles = await this.#prepareRetirement(plan.archive, 'archive');
           }
-          await this.#assertNoContextOffloadReferences?.(plan.remove.sessionIds);
           const allSessionIds = [...plan.remove.sessionIds, ...plan.archive.sessionIds];
           await this.#finalizeWorkspacePatches(allSessionIds);
           await this.#disposeBackends(allSessionIds);
@@ -672,6 +672,7 @@ export class HostSessionRetirementCoordinator {
         {
           artifacts: this.#artifacts,
           taskLedger: this.#taskLedger,
+          ...(this.#contextOffload ? { contextOffload: this.#contextOffload } : {}),
           purgeOperationalState: this.#purgeOperationalState,
         },
         sessionId,

packages/runtime-host/src/server/session-revision-coordinator.ts

@@ -35,6 +35,7 @@ import {
   archivedToolResultContainsConversationOwnedReferences,
   cloneConversationRuntimeLedger,
   collectConversationCopyLinkedChildReferences,
+  collectConversationCopySessionContextRefIds,
   createConversationCopySlice,
   prepareConversationRuntimeLedgerCopy,
   type ConversationRuntimeLedgerCopyPlan,
@@ -54,6 +55,7 @@ import {
   authenticateInteractiveTaskLedgerWriter,
   type InteractiveTaskLedgerWriter,
 } from '@maka/storage/task-ledger-authority';
+import type { InteractiveContextOffloadWriter } from '@maka/storage/context-offload-store';
 import type {
   OperationOutcome,
   SessionConversationCopyInput,
@@ -99,6 +101,10 @@ export interface HostSessionRevisionCoordinatorOptions {
   readonly stores: ExecutionStoresWriter<'interactive'>;
   readonly artifacts: InteractiveArtifactStoreWriter;
   readonly taskLedger: InteractiveTaskLedgerWriter;
+  readonly contextOffload?: Pick<
+    InteractiveContextOffloadWriter,
+    'copyReferences' | 'retireSession'
+  >;
   readonly manager: SessionManager;
   readonly admission: SessionAdmissionGate;
   readonly continuity: SessionContinuityCoordinator;
@@ -489,6 +495,28 @@ export class HostSessionRevisionCoordinator {
           )
           .map(({ descriptor, serializedResult }) => [descriptor.artifactId, serializedResult]),
       );
+      const sourceContextRefIds = collectConversationCopySessionContextRefIds({
+        sourceSessionId: input.sourceSessionId,
+        copiedMessages: slice.messages,
+        plan,
+      });
+      if (sourceContextRefIds.length > 0 && !this.options.contextOffload) {
+        throw new Error('Session context copy authority is unavailable');
+      }
+      const contextCopy =
+        sourceContextRefIds.length === 0
+          ? { ok: true as const, copied: [] }
+          : await this.options.contextOffload!.copyReferences({
+              sourceSessionId: input.sourceSessionId,
+              targetSessionId: input.targetSessionId,
+              references: sourceContextRefIds.map((sourceRefId) => ({
+                sourceRefId,
+                targetOwner: { kind: 'read_image_snapshot', ownerId: sourceRefId },
+              })),
+            });
+      if (!contextCopy.ok) {
+        throw new Error(`Session context references could not be copied: ${contextCopy.reason}`);
+      }
       const artifactCopy = await this.#artifacts.copyConversationArtifacts({
         sourceSessionId: input.sourceSessionId,
         targetSessionId: input.targetSessionId,
@@ -511,6 +539,9 @@ export class HostSessionRevisionCoordinator {
         targetSessionId: input.targetSessionId,
         artifactIds: artifactCopy.artifactIds,
         relativePaths: artifactCopy.relativePaths,
+        contextRefs: new Map(
+          contextCopy.copied.map(({ sourceRefId, targetRefId }) => [sourceRefId, targetRefId]),
+        ),
         linkedChildren:
           kind === 'side_conversation'
             ? {
@@ -803,6 +834,7 @@ export class HostSessionRevisionCoordinator {
       {
         artifacts: this.#artifacts,
         taskLedger: this.#taskLedger,
+        ...(this.options.contextOffload ? { contextOffload: this.options.contextOffload } : {}),
         purgeOperationalState: (sessionId) =>
           this.#stores.purgeConversationOperationalState(sessionId),
       },

packages/runtime-host/src/server/session-sidecar-purge.ts

@@ -19,10 +19,12 @@
 
 import type { InteractiveArtifactStoreWriter } from '@maka/storage/artifact-stores';
 import type { InteractiveTaskLedgerWriter } from '@maka/storage/task-ledger-authority';
+import type { InteractiveContextOffloadWriter } from '@maka/storage/context-offload-store';
 
 export interface SessionSidecarPurgeAuthority {
   readonly artifacts: Pick<InteractiveArtifactStoreWriter, 'purgeSessionArtifacts'>;
   readonly taskLedger: Pick<InteractiveTaskLedgerWriter, 'purgeConversationTaskLedger'>;
+  readonly contextOffload?: Pick<InteractiveContextOffloadWriter, 'retireSession'>;
   readonly purgeOperationalState: (sessionId: string) => Promise<void>;
 }
 
@@ -33,6 +35,7 @@ export async function purgeSessionSidecars(
   const outcomes = await Promise.allSettled([
     authority.artifacts.purgeSessionArtifacts(sessionId),
     authority.taskLedger.purgeConversationTaskLedger(sessionId),
+    ...(authority.contextOffload ? [authority.contextOffload.retireSession(sessionId)] : []),
     authority.purgeOperationalState(sessionId),
   ]);
   const failures = outcomes.flatMap((outcome) =>

packages/runtime/src/__tests__/builtin-tools-file-worker.test.ts

@@ -169,25 +169,30 @@ describe('builtin file tools use the sandboxed worker', () => {
   test('uses one worker read operation for image paths', async () => {
     const cwd = await temporaryDirectory('maka-file-worker-cwd-');
     const calls: FilesystemWorkerExecuteInput[] = [];
+    let snapshotOwnerId: string | undefined;
     const tools = buildBuiltinTools({
       filesystemWorker: {
         execute: async (input) => {
           calls.push(input);
           return { kind: 'read_image', base64: 'iVBORw0KGgo=', mimeType: 'image/png' };
         },
       },
-      snapshotImage: async () => ({
-        kind: 'session_file',
-        sessionId: 'session-1',
-        relativePath: 'artifact-1',
-      }),
+      snapshotImage: async (input) => {
+        snapshotOwnerId = input.ownerId;
+        return {
+          kind: 'session_context',
+          sessionId: 'session-1',
+          refId: 'context-1',
+        };
+      },
       sandboxPlatform: 'darwin',
     });
 
     await runTool(tools, 'Read', { path: 'image.png', offset: 1, limit: 1 }, cwd);
 
     assert.equal(calls.length, 1);
     assert.deepEqual(calls[0]?.operation, { kind: 'read', path: 'image.png', offset: 1, limit: 1 });
+    assert.equal(snapshotOwnerId, 'tool-Read');
   });
 
   test('serializes writes through real and symlinked cwd paths', async () => {