diff --git a/CHANGELOG.md b/CHANGELOG.md index d401e8f..d6815ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ ### Features * add ref overrides for configuration fetches ([#174](https://github.com/ubiquity-os/plugin-sdk/issues/174)) ([a915b12](https://github.com/ubiquity-os/plugin-sdk/commit/a915b12ef4218830ad44ae5cb6482a5fac38217c)) +* refresh kernel attestation before LLM calls (PR [#178](https://github.com/ubiquity-os/plugin-sdk/pull/178)) ## [3.9.0](https://github.com/ubiquity-os/plugin-sdk/compare/v3.8.4...v3.9.0) (2026-01-12) diff --git a/README.md b/README.md index 794f712..05f9851 100644 --- a/README.md +++ b/README.md @@ -21,10 +21,6 @@ The `createPlugin` function enables users to create a plugin that will run on Cl The `callLlm` function sends chat completion requests to `ai.ubq.fi` using the auth token and repository context supplied by the kernel. -### `callLlm` - -The `callLlm` function sends chat completion requests to `ai.ubq.fi` using the auth token and repository context supplied by the kernel. - ### `postComment` Use `context.commentHandler.postComment` to write or update a comment on the triggering issue or pull request. @@ -87,32 +83,6 @@ const result = await callLlm( ); ``` -## LLM Utility - -```ts -import { callLlm } from "@ubiquity-os/plugin-sdk"; - -const result = await callLlm( - { - messages: [{ role: "user", content: "Summarize this issue." }], - }, - context -); -``` - -## LLM Utility - -```ts -import { callLlm } from "@ubiquity-os/plugin-sdk"; - -const result = await callLlm( - { - messages: [{ role: "user", content: "Summarize this issue." }], - }, - context -); -``` - ## Markdown Cleaning Utility `cleanMarkdown` removes top-level HTML comments and configured HTML tags while preserving content inside fenced/indented code blocks, inline code spans, and blockquotes. diff --git a/src/configuration/schema.ts b/src/configuration/schema.ts index 80674d4..37d6c28 100644 --- a/src/configuration/schema.ts +++ b/src/configuration/schema.ts @@ -44,7 +44,8 @@ export function stringLiteralUnion(values: readonly [...T]): return T.Union(literals as never); } -const emitterType = stringLiteralUnion(emitterEventNames); +const customKernelEvents = ["kernel.plugin_error"] as const; +const emitterType = stringLiteralUnion([...emitterEventNames, ...customKernelEvents] as const); const runsOnSchema = T.Array(emitterType, { default: [] }); @@ -54,7 +55,7 @@ const pluginSettingsSchema = T.Union( T.Null(), T.Object( { - with: T.Optional(T.Record(T.String(), T.Unknown(), { default: {} })), + with: T.Record(T.String(), T.Unknown(), { default: {} }), runsOn: T.Optional(runsOnSchema), skipBotEvents: T.Optional(T.Boolean()), }, diff --git a/src/llm/index.ts b/src/llm/index.ts index 0df816d..d1d63c4 100644 --- a/src/llm/index.ts +++ b/src/llm/index.ts @@ -50,6 +50,34 @@ function isGitHubToken(token: string): boolean { return token.trim().startsWith("gh"); } +type KernelRefreshTokens = { + authToken: string; + kernelToken: string; + expiresAt?: string | null; +}; + +function readKernelRefreshUrl(value: unknown): string { + if (!value || typeof value !== "object") return EMPTY_STRING; + return normalizeToken((value as Record).kernelRefreshUrl); +} + +function getKernelRefreshUrl(input: PluginInput | Context): string { + const direct = readKernelRefreshUrl(input); + if (direct) return direct; + const fromConfig = readKernelRefreshUrl((input as { config?: unknown }).config); + if (fromConfig) return fromConfig; + return readKernelRefreshUrl((input as { settings?: unknown }).settings); +} + +function updateInputTokens(input: PluginInput | Context, tokens: KernelRefreshTokens) { + if ("authToken" in input) { + (input as { authToken?: string }).authToken = tokens.authToken; + } + if ("ubiquityKernelToken" in input) { + (input as { ubiquityKernelToken?: string }).ubiquityKernelToken = tokens.kernelToken; + } +} + function getEnvTokenFromInput(input: PluginInput | Context): string { if ("env" in input) { const envValue = (input as Context).env; @@ -91,10 +119,27 @@ function getAiBaseUrl(options: LlmCallOptions): string { export async function callLlm(options: LlmCallOptions, input: PluginInput | Context): Promise> { const { baseUrl, model, stream: isStream, messages, aiAuthToken, ...rest } = options; - const { token: authToken, isGitHub } = resolveAuthToken(input, aiAuthToken); - const kernelToken = "ubiquityKernelToken" in input ? input.ubiquityKernelToken : undefined; + const { token: resolvedAuthToken, isGitHub } = resolveAuthToken(input, aiAuthToken); + let authToken = resolvedAuthToken; + let kernelToken = "ubiquityKernelToken" in input ? input.ubiquityKernelToken : undefined; const payload = getPayload(input); const { owner, repo, installationId } = getRepoMetadata(payload); + const kernelRefreshUrl = getKernelRefreshUrl(input); + const hasExplicitAiAuth = Boolean(normalizeToken(aiAuthToken)); + + if (isGitHub && kernelRefreshUrl && kernelToken && !hasExplicitAiAuth) { + const refreshed = await refreshKernelTokens({ + url: kernelRefreshUrl, + authToken, + kernelToken, + owner, + repo, + installationId, + }); + authToken = refreshed.authToken; + kernelToken = refreshed.kernelToken; + updateInputTokens(input, refreshed); + } ensureMessages(messages); const url = buildAiUrl(options, baseUrl); @@ -145,6 +190,53 @@ function buildAiUrl(options: LlmCallOptions, baseUrl?: string): string { return `${getAiBaseUrl({ ...options, baseUrl })}/v1/chat/completions`; } +async function refreshKernelTokens(params: { + url: string; + authToken: string; + kernelToken: string; + owner: string; + repo: string; + installationId?: number; +}): Promise { + const headers = buildHeaders(params.authToken, { + owner: params.owner, + repo: params.repo, + installationId: params.installationId, + ubiquityKernelToken: params.kernelToken, + }); + const response = await fetch(params.url, { method: "POST", headers }); + const text = await response.text().catch(() => EMPTY_STRING); + if (!response.ok) { + const error = new Error(`Kernel refresh error: ${response.status} - ${text}`); + (error as Error & { status?: number }).status = response.status; + throw error; + } + + let payload: Record = {}; + if (text.trim()) { + try { + payload = JSON.parse(text) as Record; + } catch (err) { + const details = err instanceof Error ? ` (${err.message})` : ""; + const error = new Error(`Kernel refresh error: failed to parse JSON response${details}`); + (error as Error & { status?: number }).status = response.status; + throw error; + } + } + + const authToken = normalizeToken(payload.authToken); + const kernelToken = normalizeToken(payload.ubiquityKernelToken); + if (!authToken || !kernelToken) { + throw new Error("Kernel refresh error: response missing authToken or ubiquityKernelToken"); + } + + return { + authToken, + kernelToken, + expiresAt: typeof payload.expiresAt === "string" ? payload.expiresAt : null, + }; +} + async function fetchWithRetry(url: string, options: RequestInit, maxRetries: number): Promise { let attempt = 0; let lastError: unknown; diff --git a/src/types/manifest.ts b/src/types/manifest.ts index 36e2670..b2a4d4d 100644 --- a/src/types/manifest.ts +++ b/src/types/manifest.ts @@ -1,7 +1,8 @@ import { type Static, Type as T } from "@sinclair/typebox"; import { emitterEventNames } from "@octokit/webhooks"; -export const runEvent = T.Union(emitterEventNames.map((o) => T.Literal(o))); +const customKernelEvents = ["kernel.plugin_error"] as const; +export const runEvent = T.Union([...emitterEventNames, ...customKernelEvents].map((o) => T.Literal(o))); export const exampleCommandExecutionSchema = T.Object({ commandInvocation: T.String({ minLength: 1 }), diff --git a/tests/llm.test.ts b/tests/llm.test.ts index e742a08..9541468 100644 --- a/tests/llm.test.ts +++ b/tests/llm.test.ts @@ -125,6 +125,71 @@ describe("callLlm", () => { ); }); + it("refreshes kernel attestation before calling the LLM for GitHub auth", async () => { + const completion = { + id: "completion-2", + object: "chat.completion", + created: 1, + model: "gpt-5.1", + choices: [], + } as ChatCompletion; + + const input = { + authToken: "ghs_initial_token", + ubiquityKernelToken: "kernel-initial", + eventPayload: { + repository: { owner: { login: "octo" }, name: "repo" }, + installation: { id: 123 }, + }, + config: { + kernelRefreshUrl: "https://kernel.test/internal/agent/refresh-token", + }, + }; + + const fetchMock = jest.spyOn(globalThis, "fetch").mockImplementation(async (url) => { + if (String(url) === "https://kernel.test/internal/agent/refresh-token") { + return new Response(JSON.stringify({ authToken: "ghs_refreshed", ubiquityKernelToken: "kernel-refreshed" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (String(url) === "https://ai.ubq.fi/v1/chat/completions") { + return new Response(JSON.stringify(completion), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("not found", { status: 404 }); + }); + + const result = await callLlm({ messages: [{ role: "user", content: "Hi" }], baseUrl: "https://ai.ubq.fi" }, input as typeof baseInput); + + expect(result).toEqual(completion); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(fetchMock).toHaveBeenCalledWith( + "https://kernel.test/internal/agent/refresh-token", + expect.objectContaining({ + method: "POST", + headers: expect.objectContaining({ + Authorization: "Bearer ghs_initial_token", + "X-Ubiquity-Kernel-Token": "kernel-initial", + "X-GitHub-Owner": "octo", + "X-GitHub-Repo": "repo", + "X-GitHub-Installation-Id": "123", + }), + }) + ); + expect(fetchMock).toHaveBeenCalledWith( + "https://ai.ubq.fi/v1/chat/completions", + expect.objectContaining({ + headers: expect.objectContaining({ + Authorization: "Bearer ghs_refreshed", + "X-Ubiquity-Kernel-Token": "kernel-refreshed", + }), + }) + ); + }); + it("throws on API errors", async () => { const fetchMock = jest.spyOn(globalThis, "fetch").mockResolvedValue(new Response("bad request", { status: 400 }));