Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions apps/mobile/src/lib/modelOptions.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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([]);
});
});
9 changes: 8 additions & 1 deletion apps/mobile/src/lib/modelOptions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
3 changes: 3 additions & 0 deletions apps/server/src/atomicWrite.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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* () {
Expand All @@ -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);
}),
);
1 change: 1 addition & 0 deletions apps/server/src/auth/RpcAuthorization.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
28 changes: 25 additions & 3 deletions apps/server/src/mcp/McpProviderSession.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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();
}
22 changes: 22 additions & 0 deletions apps/server/src/mcp/McpSessionRegistry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}),
);
98 changes: 74 additions & 24 deletions apps/server/src/mcp/McpSessionRegistry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand Down Expand Up @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -188,6 +233,7 @@ const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* (

return McpSessionRegistry.of({
issue,
issueIfCurrent,
resolve,
touch,
revokeProviderSession: Effect.fn("McpSessionRegistry.revokeProviderSession")(
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading