feat(prime): isolate native instance ownership

Refs #199

Author: rynfar

apps/mobile/src/lib/modelOptions.test.ts

@@ -432,4 +432,27 @@ describe("mobile model options", () => {
     expect(resolve(null, null, sticky)).toBe(sticky);
     expect(resolve(null, null, null)).toBe(providerDefault.selection);
   });
+
+  it("omits a disabled fallback selection from mobile options", () => {
+    const fallback = {
+      instanceId: ProviderInstanceId.make("primeAgent"),
+      model: "default",
+    };
+    const config = {
+      providers: [
+        {
+          instanceId: "primeAgent",
+          driver: "primeAgent",
+          displayName: "Prime Agent",
+          enabled: false,
+          installed: true,
+          status: "disabled",
+          auth: { status: "authenticated" },
+          models: [],
+        },
+      ],
+    } as unknown as ServerConfig;
+
+    expect(buildModelOptions(config, fallback)).toEqual([]);
+  });
 });

apps/mobile/src/lib/modelOptions.ts

@@ -297,7 +297,14 @@ export function buildModelOptions(
       const provider = config?.providers.find(
         (candidate) => candidate.instanceId === fallbackModelSelection.instanceId,
       );
-      if (getProviderUnavailablePresentation(provider) === null) {
+      if (
+        provider !== undefined &&
+        getProviderAdmissionAvailability({
+          provider,
+          instanceId: String(fallbackModelSelection.instanceId),
+          providerSnapshotKnown: true,
+        }).status === "available"
+      ) {
         const providerLabel = provider
           ? providerDisplayLabel(provider)
           : fallbackModelSelection.instanceId;

apps/server/src/atomicWrite.ts

@@ -5,6 +5,8 @@ import * as Path from "effect/Path";
 export const writeFileStringAtomically = (input: {
   readonly filePath: string;
   readonly contents: string;
+  /** Optional process-local fence checked immediately before the atomic rename. */
+  readonly commitGuard?: Effect.Effect<boolean>;
 }) =>
   Effect.scoped(
     Effect.gen(function* () {
@@ -20,6 +22,7 @@ export const writeFileStringAtomically = (input: {
       const tempPath = path.join(tempDirectory, "contents.tmp");
 
       yield* fs.writeFileString(tempPath, input.contents);
+      if (input.commitGuard !== undefined && !(yield* input.commitGuard)) return;
       yield* fs.rename(tempPath, input.filePath);
     }),
   );

apps/server/src/auth/RpcAuthorization.ts

@@ -62,6 +62,7 @@ export const RPC_REQUIRED_SCOPES = {
   [WS_METHODS.serverRemoveKeybinding]: AuthOrchestrationOperateScope,
   [WS_METHODS.serverGetSettings]: AuthOrchestrationReadScope,
   [WS_METHODS.serverUpdateSettings]: AuthOrchestrationOperateScope,
+  [WS_METHODS.serverMutateProviderInstances]: AuthOrchestrationOperateScope,
   [WS_METHODS.serverDiscoverSourceControl]: AuthOrchestrationReadScope,
   [WS_METHODS.serverGetTraceDiagnostics]: AuthOrchestrationReadScope,
   [WS_METHODS.serverGetProcessDiagnostics]: AuthOrchestrationReadScope,

apps/server/src/mcp/McpProviderSession.ts

@@ -1,4 +1,5 @@
 import type { EnvironmentId, ProviderInstanceId, ThreadId } from "@t3tools/contracts";
+import type { ProviderRuntimeFence } from "../provider/ProviderDriver.ts";
 
 export interface McpProviderSessionConfig {
   readonly environmentId: EnvironmentId;
@@ -10,19 +11,40 @@ export interface McpProviderSessionConfig {
 }
 
 const sessionsByThread = new Map<ThreadId, McpProviderSessionConfig>();
+const generationsByThread = new Map<ThreadId, object>();
 
-export function setMcpProviderSession(config: McpProviderSessionConfig): void {
+export function setMcpProviderSession(
+  config: McpProviderSessionConfig,
+  runtimeFence?: ProviderRuntimeFence,
+): void {
   sessionsByThread.set(config.threadId, config);
+  if (runtimeFence === undefined) generationsByThread.delete(config.threadId);
+  else generationsByThread.set(config.threadId, runtimeFence.generation);
 }
 
 export function readMcpProviderSession(threadId: ThreadId): McpProviderSessionConfig | undefined {
   return sessionsByThread.get(threadId);
 }
 
-export function clearMcpProviderSession(threadId: ThreadId): void {
-  sessionsByThread.delete(threadId);
+export function isMcpProviderSessionOwnedByGeneration(
+  threadId: ThreadId,
+  runtimeFence: ProviderRuntimeFence,
+): boolean {
+  return generationsByThread.get(threadId) === runtimeFence.generation;
+}
+
+export function clearMcpProviderSession(
+  threadId: ThreadId,
+  runtimeFence?: ProviderRuntimeFence,
+): boolean {
+  if (runtimeFence !== undefined && generationsByThread.get(threadId) !== runtimeFence.generation) {
+    return false;
+  }
+  generationsByThread.delete(threadId);
+  return sessionsByThread.delete(threadId);
 }
 
 export function clearAllMcpProviderSessions(): void {
   sessionsByThread.clear();
+  generationsByThread.clear();
 }

apps/server/src/mcp/McpSessionRegistry.test.ts

@@ -127,3 +127,25 @@ it.effect("does not keep credentials of other threads alive", () =>
     expect(yield* registry.resolve(token)).toBeUndefined();
   }),
 );
+
+it.effect("keeps the current exact credential when retired issue and cleanup arrive late", () =>
+  Effect.gen(function* () {
+    const registry = yield* makeRegistry(() => 1_000);
+    const threadId = ThreadId.make("thread-generation-fence");
+    const request = {
+      threadId,
+      providerInstanceId: ProviderInstanceId.make("primeAgent"),
+    };
+    const first = yield* registry.issue(request);
+    const firstToken = first.config.authorizationHeader.replace(/^Bearer\s+/, "");
+    const replacement = yield* registry.issueIfCurrent(request, Effect.succeed(true));
+    expect(replacement).toBeDefined();
+    const replacementToken = replacement!.config.authorizationHeader.replace(/^Bearer\s+/, "");
+    expect(yield* registry.resolve(firstToken)).toBeUndefined();
+
+    const retiredIssue = yield* registry.issueIfCurrent(request, Effect.succeed(false));
+    expect(retiredIssue).toBeUndefined();
+    yield* registry.revokeProviderSession(first.config.providerSessionId);
+    expect((yield* registry.resolve(replacementToken))?.threadId).toBe(threadId);
+  }),
+);

apps/server/src/mcp/McpSessionRegistry.ts

@@ -22,6 +22,11 @@ export interface McpIssuedCredential {
 
 export interface McpSessionRegistryShape {
   readonly issue: (request: McpCredentialRequest) => Effect.Effect<McpIssuedCredential>;
+  /** Atomically replace one thread credential only while its provider generation is current. */
+  readonly issueIfCurrent: (
+    request: McpCredentialRequest,
+    isCurrent: Effect.Effect<boolean>,
+  ) => Effect.Effect<McpIssuedCredential | undefined>;
   readonly resolve: (
     rawToken: string,
   ) => Effect.Effect<McpInvocationContext.McpInvocationScope | undefined>;
@@ -117,26 +122,26 @@ const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* (
     return next.size === records.size ? records : next;
   };
 
-  const issue: McpSessionRegistryShape["issue"] = Effect.fn("McpSessionRegistry.issue")(
-    function* (request) {
-      const issuedAt = yield* currentTimeMillis;
-      const providerSessionId = yield* crypto.randomUUIDv4.pipe(Effect.orDie);
-      const rawToken = yield* crypto.randomBytes(32).pipe(Effect.map(tokenFromBytes), Effect.orDie);
-      const tokenHash = yield* hashToken(rawToken);
-      const scope: McpInvocationContext.McpInvocationScope = {
-        environmentId,
-        threadId: ThreadId.make(request.threadId),
-        providerSessionId,
-        providerInstanceId: ProviderInstanceId.make(request.providerInstanceId),
-        capabilities: new Set(["preview"]),
-        issuedAt,
-      };
-      yield* SynchronizedRef.update(state, ({ records }) => {
-        const next = new Map(pruneDead(records, issuedAt));
-        next.set(tokenHash, { tokenHash, scope, lastAliveAt: issuedAt });
-        return { records: next };
-      });
-      return {
+  const prepareCredential = Effect.fn("McpSessionRegistry.prepareCredential")(function* (
+    request: McpCredentialRequest,
+  ) {
+    const issuedAt = yield* currentTimeMillis;
+    const providerSessionId = yield* crypto.randomUUIDv4.pipe(Effect.orDie);
+    const rawToken = yield* crypto.randomBytes(32).pipe(Effect.map(tokenFromBytes), Effect.orDie);
+    const tokenHash = yield* hashToken(rawToken);
+    const scope: McpInvocationContext.McpInvocationScope = {
+      environmentId,
+      threadId: ThreadId.make(request.threadId),
+      providerSessionId,
+      providerInstanceId: ProviderInstanceId.make(request.providerInstanceId),
+      capabilities: new Set(["preview"]),
+      issuedAt,
+    };
+    return {
+      issuedAt,
+      tokenHash,
+      scope,
+      credential: {
         config: {
           environmentId,
           threadId: scope.threadId,
@@ -145,10 +150,50 @@ const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* (
           endpoint,
           authorizationHeader: `Bearer ${rawToken}`,
         },
-      };
+      } satisfies McpIssuedCredential,
+    };
+  });
+
+  const issue: McpSessionRegistryShape["issue"] = Effect.fn("McpSessionRegistry.issue")(
+    function* (request) {
+      const prepared = yield* prepareCredential(request);
+      yield* SynchronizedRef.update(state, ({ records }) => {
+        const next = new Map(pruneDead(records, prepared.issuedAt));
+        next.set(prepared.tokenHash, {
+          tokenHash: prepared.tokenHash,
+          scope: prepared.scope,
+          lastAliveAt: prepared.issuedAt,
+        });
+        return { records: next };
+      });
+      return prepared.credential;
     },
   );
 
+  const issueIfCurrent: McpSessionRegistryShape["issueIfCurrent"] = Effect.fn(
+    "McpSessionRegistry.issueIfCurrent",
+  )(function* (request, isCurrent) {
+    const prepared = yield* prepareCredential(request);
+    return yield* SynchronizedRef.modifyEffect(state, ({ records }) =>
+      Effect.gen(function* () {
+        // The generation check and replacement share the registry's single mutation permit.
+        if (!(yield* isCurrent)) return [undefined, { records }] as const;
+        const current = pruneDead(records, prepared.issuedAt);
+        const next = new Map(
+          Array.from(current).filter(
+            ([, record]) => record.scope.threadId !== prepared.scope.threadId,
+          ),
+        );
+        next.set(prepared.tokenHash, {
+          tokenHash: prepared.tokenHash,
+          scope: prepared.scope,
+          lastAliveAt: prepared.issuedAt,
+        });
+        return [prepared.credential, { records: next }] as const;
+      }),
+    );
+  });
+
   const resolve: McpSessionRegistryShape["resolve"] = Effect.fn("McpSessionRegistry.resolve")(
     function* (rawToken) {
       if (rawToken.length === 0) return undefined;
@@ -188,6 +233,7 @@ const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* (
 
   return McpSessionRegistry.of({
     issue,
+    issueIfCurrent,
     resolve,
     touch,
     revokeProviderSession: Effect.fn("McpSessionRegistry.revokeProviderSession")(
@@ -224,13 +270,17 @@ export const layer = Layer.effect(McpSessionRegistry, make);
 
 export const issueActiveMcpCredential = (
   request: McpCredentialRequest,
+  isCurrent: Effect.Effect<boolean> = Effect.succeed(true),
 ): Effect.Effect<McpIssuedCredential | undefined> =>
   activeMcpSessionRegistry
-    ? activeMcpSessionRegistry
-        .revokeThread(request.threadId)
-        .pipe(Effect.andThen(activeMcpSessionRegistry.issue(request)))
+    ? activeMcpSessionRegistry.issueIfCurrent(request, isCurrent)
     : Effect.sync((): McpIssuedCredential | undefined => undefined);
 
+export const revokeActiveMcpProviderSession = (providerSessionId: string): Effect.Effect<void> =>
+  activeMcpSessionRegistry
+    ? activeMcpSessionRegistry.revokeProviderSession(providerSessionId)
+    : Effect.void;
+
 /**
  * Refreshes the liveness of a thread's MCP credential. Called on every provider
  * turn so an active session is never mistaken for an abandoned one.