Skip to content

fix(pi): resolve native embedding routing before registration #462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ualtinok merged 2 commits into from
Sep 17, 2026
+223 −11
Merged

fix(pi): resolve native embedding routing before registration #462

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ad239a0
fix(pi): resolve native embedding routing before registration
Qiiks Sep 17, 2026
c972f50
fix(pi): retire stale shadow routing and memoize static fallback
Qiiks Sep 17, 2026
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
158 changes: 156 additions & 2 deletions packages/pi-plugin/src/embedding-bootstrap.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,21 @@
import { describe, expect, it } from "bun:test";
import { getProjectEmbeddingSnapshot } from "@magic-context/core/features/magic-context/memory/embedding";
import { describe, expect, it, spyOn } from "bun:test";
import { promises as fs } from "node:fs";
import path from "node:path";
import type { EmbeddingConfig } from "@magic-context/core/config/schema/magic-context";
import {
_resetProjectEmbeddingRegistryForTests,
_setTestProviderFactoryForProject,
getProjectEmbeddingSnapshot,
getShadowEmbeddingMeasurementCohort,
registerProjectShadowEmbedding,
} from "@magic-context/core/features/magic-context/memory/embedding";
import {
getProjectEmbeddings,
peekProjectEmbeddings,
resetEmbeddingCacheForTests,
} from "@magic-context/core/features/magic-context/memory/embedding-cache";
import { resolveProjectIdentity } from "@magic-context/core/features/magic-context/memory/project-identity";
import * as logger from "@magic-context/core/shared/logger";
import { closeQuietly } from "@magic-context/core/shared/sqlite-helpers";
import { createTestTempDir } from "@magic-context/core/shared/test-temp-dir";

Expand All @@ -16,9 +26,11 @@ describe("ensureProjectRegisteredFromPiDirectory", () => {
it("preserves the embedding cache across consecutive identical registrations", async () => {
const db = createTestDb();
const oldHome = process.env.HOME;
const oldConfigHome = process.env.XDG_CONFIG_HOME;
const directory = createTestTempDir("pi-embedding-bootstrap-").dir;
const fakeHome = createTestTempDir("pi-embedding-home-").dir;
process.env.HOME = fakeHome;
process.env.XDG_CONFIG_HOME = path.join(fakeHome, ".config");
resetEmbeddingCacheForTests();
try {
const projectIdentity = resolveProjectIdentity(directory);
Expand All @@ -43,6 +55,148 @@ describe("ensureProjectRegisteredFromPiDirectory", () => {
} else {
process.env.HOME = oldHome;
}
if (oldConfigHome === undefined) delete process.env.XDG_CONFIG_HOME;
else process.env.XDG_CONFIG_HOME = oldConfigHome;
closeQuietly(db);
}
});

it("registers the fallback identity when native discovery is unavailable", async () => {
const db = createTestDb();
const directory = createTestTempDir("pi-embedding-synapse-").dir;
const fakeHome = createTestTempDir("pi-embedding-synapse-home-").dir;
const previous = {
HOME: process.env.HOME,
XDG_CONFIG_HOME: process.env.XDG_CONFIG_HOME,
};
process.env.HOME = fakeHome;
process.env.XDG_CONFIG_HOME = path.join(fakeHome, ".config");
resetEmbeddingCacheForTests();
try {
// Provider and SubC settings are user-tier only.
const configDir = path.join(fakeHome, ".config", "cortexkit");
await fs.mkdir(configDir, { recursive: true });
await fs.writeFile(
path.join(configDir, "magic-context.json"),
JSON.stringify({
embedding: { provider: "synapse", fallback_provider: "off" },
subc: { connection_file: path.join(fakeHome, "absent-subc.json") },
}),
);
const projectIdentity = resolveProjectIdentity(directory);
await ensureProjectRegisteredFromPiDirectory(directory, db);
expect(getProjectEmbeddingSnapshot(projectIdentity)?.provider).toBe(
"off",
);
expect(getProjectEmbeddingSnapshot(projectIdentity)?.modelId).not.toMatch(
/synapse/u,
);
} finally {
resetEmbeddingCacheForTests();
for (const [key, value] of Object.entries(previous)) {
if (value === undefined) delete process.env[key];
else process.env[key] = value;
}
closeQuietly(db);
}
});
it("retires a disabled shadow without removing the primary lane", async () => {
const db = createTestDb();
const directory = createTestTempDir("pi-shadow-retirement-").dir;
const configHome = createTestTempDir("pi-shadow-config-").dir;
const previous = process.env.XDG_CONFIG_HOME;
process.env.XDG_CONFIG_HOME = configHome;
let disposed = false;
_setTestProviderFactoryForProject(() => ({
modelId: "shadow",
initialize: async () => true,
embed: async () => new Float32Array([1, 0]),
embedBatch: async (texts: string[]) =>
texts.map(() => new Float32Array([1, 0])),
dispose: async () => {
disposed = true;
},
isLoaded: () => true,
}));
try {
await fs.mkdir(path.join(configHome, "cortexkit"), { recursive: true });
await fs.writeFile(
path.join(configHome, "cortexkit", "magic-context.json"),
JSON.stringify({
embedding: { provider: "off" },
shadow_embedding: { enabled: false },
}),
);
const identity = resolveProjectIdentity(directory);
registerProjectShadowEmbedding(
db,
identity,
{
provider: "synapse",
model: "shadow",
synapse_fingerprint: "fixture",
} as unknown as EmbeddingConfig,
directory,
);
expect(getShadowEmbeddingMeasurementCohort(identity)?.fingerprint).toBe(
"fixture",
);
await ensureProjectRegisteredFromPiDirectory(directory, db);
expect(getShadowEmbeddingMeasurementCohort(identity)).toBeNull();
expect(getProjectEmbeddingSnapshot(identity)?.provider).toBe("off");
expect(disposed).toBe(true);

@cubic-dev-ai cubic-dev-ai Bot Sep 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3: The expect(disposed).toBe(true) assertion depends on a synchronous timing detail: unregisterProjectShadowEmbedding → disposeProvider calls void provider.dispose() without awaiting, and the test provider's async dispose body sets disposed = true before its first (nonexistent) await. This passes only because the async body runs synchronously. If disposeProvider ever starts deferring disposal or the provider discards returns control before setting the flag, the test would silently pass/fail based on timing rather than actual retirement. Since the disposal is fire-and-forget in the production code, assert on the deterministic observable (shadow cohort null, primary intact) and avoid relying on the fire-and-forget side effect timing for the disposal check.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At packages/pi-plugin/src/embedding-bootstrap.test.ts, line 147:

<comment>The `expect(disposed).toBe(true)` assertion depends on a synchronous timing detail: `unregisterProjectShadowEmbedding` → `disposeProvider` calls `void provider.dispose()` without awaiting, and the test provider's async `dispose` body sets `disposed = true` before its first (nonexistent) await. This passes only because the async body runs synchronously. If `disposeProvider` ever starts deferring disposal or the provider discards returns control before setting the flag, the test would silently pass/fail based on timing rather than actual retirement. Since the disposal is fire-and-forget in the production code, assert on the deterministic observable (shadow cohort null, primary intact) and avoid relying on the fire-and-forget side effect timing for the disposal check.</comment>

<file context>
@@ -92,4 +100,104 @@ describe("ensureProjectRegisteredFromPiDirectory", () => {
+			await ensureProjectRegisteredFromPiDirectory(directory, db);
+			expect(getShadowEmbeddingMeasurementCohort(identity)).toBeNull();
+			expect(getProjectEmbeddingSnapshot(identity)?.provider).toBe("off");
+			expect(disposed).toBe(true);
+		} finally {
+			_resetProjectEmbeddingRegistryForTests();
</file context>

} finally {
_resetProjectEmbeddingRegistryForTests();
if (previous === undefined) delete process.env.XDG_CONFIG_HOME;
else process.env.XDG_CONFIG_HOME = previous;
closeQuietly(db);
}
});
it("does not repeat a missing-SubC warning until configuration changes", async () => {
const db = createTestDb();
const directory = createTestTempDir("pi-routing-memo-").dir;
const configHome = createTestTempDir("pi-routing-memo-config-").dir;
const previous = process.env.XDG_CONFIG_HOME;
process.env.XDG_CONFIG_HOME = configHome;
const messages: string[] = [];
const logging = spyOn(logger, "log").mockImplementation((message) => {
messages.push(String(message));
});
try {
const configDir = path.join(configHome, "cortexkit");
await fs.mkdir(configDir, { recursive: true });
const configFile = path.join(configDir, "magic-context.json");
await fs.writeFile(
configFile,
JSON.stringify({
embedding: { provider: "synapse", fallback_provider: "off" },
}),
);
await ensureProjectRegisteredFromPiDirectory(directory, db);
await ensureProjectRegisteredFromPiDirectory(directory, db);
expect(
messages.filter((message) => message.includes("requires a subc block")),
).toHaveLength(1);
await fs.writeFile(
configFile,
JSON.stringify({
embedding: { provider: "synapse", fallback_provider: "off" },
subc: { connection_file: path.join(configHome, "missing.json") },
}),
);
await ensureProjectRegisteredFromPiDirectory(directory, db);
await ensureProjectRegisteredFromPiDirectory(directory, db);
// A configured but unavailable daemon is retryable, unlike missing configuration.
expect(
messages.filter((message) =>
message.startsWith("[magic-context] Synapse is not ready;"),
),
).toHaveLength(2);
} finally {
logging.mockRestore();
_resetProjectEmbeddingRegistryForTests();
if (previous === undefined) delete process.env.XDG_CONFIG_HOME;
else process.env.XDG_CONFIG_HOME = previous;
closeQuietly(db);
}
});
Expand Down
58 changes: 49 additions & 9 deletions packages/pi-plugin/src/embedding-bootstrap.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,17 @@ import {
import {
type EmbeddingFeatures,
registerProjectEmbedding,
registerProjectShadowEmbedding,
unregisterProjectShadowEmbedding,
} from "@magic-context/core/features/magic-context/memory/embedding";
import { resolveProjectIdentityForSession } from "@magic-context/core/features/magic-context/memory/project-identity";
import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage";
import {
handleUntrustedLoad,
isConfigLoadUntrusted,
} from "@magic-context/core/plugin/embedding-bootstrap-helpers";
import { resolveEmbeddingRouting } from "@magic-context/core/plugin/embedding-routing";
import { log } from "@magic-context/core/shared/logger";
import { loadPiConfigDetailed } from "./config";

interface RegistrationFingerprint {
Expand Down Expand Up @@ -76,23 +80,59 @@ export async function ensureProjectRegisteredFromPiDirectory(
return;
}

const routing = await resolveEmbeddingRouting({
config: detailed.config,
projectRoot: directory,
session: `bootstrap:${projectIdentity}`,
});
for (const warning of routing.warnings) {
log(`[magic-context] ${warning}`);
}

const features: EmbeddingFeatures = {
memoryEnabled: detailed.config.memory.enabled,
gitCommitEnabled: detailed.config.memory.git_commit_indexing.enabled,
};
registerProjectEmbedding(
db,
projectIdentity,
detailed.config.embedding,
routing.primary,
features,
directory,
);
const fingerprintPaths = configCandidatePaths(
directory,
detailed.loadedFromPaths,
);
registrationFingerprints.set(projectIdentity, {
paths: fingerprintPaths,
fingerprint: configFingerprint(fingerprintPaths),
});
if (routing.shadow) {
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
registerProjectShadowEmbedding(
db,
projectIdentity,
routing.shadow,
directory,
);
} else {
unregisterProjectShadowEmbedding(projectIdentity);
}
Comment thread
greptile-apps[bot] marked this conversation as resolved.
// Only failed daemon discovery can recover without a configuration change.
const configuredProvider = detailed.config.embedding.provider;
const canDiscover =
Boolean(detailed.config.subc) &&
(configuredProvider === "synapse"
? Boolean(detailed.config.embedding.fallback_provider)
: configuredProvider !== "off" &&
detailed.config.shadow_embedding?.enabled === true);
const discoveryFailed =
canDiscover &&
(configuredProvider === "synapse"
? routing.primary.provider !== "synapse"
: routing.shadow === null);
if (!discoveryFailed) {
const fingerprintPaths = configCandidatePaths(
directory,
detailed.loadedFromPaths,
);
registrationFingerprints.set(projectIdentity, {
paths: fingerprintPaths,
fingerprint: configFingerprint(fingerprintPaths),
});
} else {
registrationFingerprints.delete(projectIdentity);
}
}
1 change: 1 addition & 0 deletions packages/plugin/src/features/magic-context/memory/embedding.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ export {
type ShadowEmbeddingMeasurementCohort,
sweepAllRegisteredProjects,
unregisterProjectEmbedding,
unregisterProjectShadowEmbedding,
} from "../project-embedding-registry";

const DEFAULT_EMBEDDING_CONFIG: EmbeddingConfig = {
Expand Down
17 changes: 17 additions & 0 deletions packages/plugin/src/features/magic-context/project-embedding-registry.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2232,6 +2232,23 @@ export function registerProjectInObservationMode(
return snapshotFor(registration);
}

export function unregisterProjectShadowEmbedding(projectIdentity: string): void {
const shadow = shadowRegistrations.get(projectIdentity);
shadowRegistrations.delete(projectIdentity);
dbForShadowQueue.delete(projectIdentity);
pendingShadowBackfills.delete(projectIdentity);
for (let index = shadowQueue.length - 1; index >= 0; index -= 1) {
if (shadowQueue[index].projectIdentity === projectIdentity) shadowQueue.splice(index, 1);
}
for (const scope of ["memory", "commit", "chunk"] as const) {
const key = `${projectIdentity}:${scope}`;
shadowBackfillLastIds.delete(key);
shadowBackfillStopReasons.delete(key);
shadowBackfillLastWriteOutcomes.delete(key);
}
disposeProvider(shadow?.provider ?? null);

@cubic-dev-ai cubic-dev-ai Bot Sep 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: If unregister runs while a shadow batch is in flight, the worker can still write vectors for the retired shadow model after this disposal. Add cancellation or a registration-generation check before committing the batch, and suppress stale worker state updates.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At packages/plugin/src/features/magic-context/project-embedding-registry.ts, line 2249:

<comment>If unregister runs while a shadow batch is in flight, the worker can still write vectors for the retired shadow model after this disposal. Add cancellation or a registration-generation check before committing the batch, and suppress stale worker state updates.</comment>

<file context>
@@ -2232,6 +2232,23 @@ export function registerProjectInObservationMode(
+        shadowBackfillStopReasons.delete(key);
+        shadowBackfillLastWriteOutcomes.delete(key);
+    }
+    disposeProvider(shadow?.provider ?? null);
+}
+
</file context>

}

export function unregisterProjectEmbedding(projectIdentity: string): void {
const prior = projectRegistrations.get(projectIdentity);
const shadow = shadowRegistrations.get(projectIdentity);
Expand Down