Skip to content
Closed
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
55 changes: 27 additions & 28 deletions apps/playground/src/app/api/chat/route.ts
Original file line number Diff line number Diff line change
@@ -1,28 +1,29 @@
// import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import {
streamText,
type UIMessage,
convertToModelMessages,
// experimental_createMCPClient,
experimental_createMCPClient,
} from "ai";

import { getUser } from "@/lib/getUser";

import { createLLMGateway } from "@llmgateway/ai-sdk-provider";

import type { LLMGatewayChatModelId } from "@llmgateway/ai-sdk-provider/internal";
// import type { experimental_MCPClient } from "ai";
import type { experimental_MCPClient } from "ai";

export const maxDuration = 300; // 5 minutes

interface ChatRequestBody {
messages: UIMessage[];
model?: LLMGatewayChatModelId;
apiKey?: string;
githubToken?: string;
}

// let githubMCP: experimental_MCPClient | null = null;
// let tools: any | null = null;
let githubMCP: experimental_MCPClient | null = null;
let tools: any | null = null;

export async function POST(req: Request) {
const user = await getUser();
Expand All @@ -34,25 +35,26 @@ export async function POST(req: Request) {
}

const body = await req.json();
const { messages, model, apiKey }: ChatRequestBody = body;

// if (!githubMCP) {
// const transport = new StreamableHTTPClientTransport(
// new URL("https://api.githubcopilot.com/mcp"),
// {
// requestInit: {
// method: "POST",
// headers: {
// Authorization: `Bearer ${githubToken}`,
// },
// },
// },
// );
// githubMCP = await experimental_createMCPClient({ transport });
// if (!tools) {
// tools = await githubMCP.tools();
// }
// }
const { messages, model, apiKey, githubToken }: ChatRequestBody = body;

if (!githubMCP && githubToken) {
const transport = new StreamableHTTPClientTransport(
new URL("https://api.githubcopilot.com/mcp"),
{
requestInit: {
method: "POST",
headers: {
Authorization: `Bearer ${githubToken}`,
},
},
},
);
githubMCP = await experimental_createMCPClient({ transport });

if (!tools) {
tools = await githubMCP.tools();
}
}

if (!messages || !Array.isArray(messages)) {
return new Response(JSON.stringify({ error: "Missing messages" }), {
Expand Down Expand Up @@ -81,9 +83,6 @@ export async function POST(req: Request) {
baseUrl: gatewayUrl,
headers: {
"x-source": "chat.llmgateway.io",
"X-LLMGateway-User-ID": user.id,
"X-LLMGateway-User-Email": user.email,
"X-LLMGateway-User-Name": user.name,
},
});
const selectedModel = (model ??
Expand All @@ -94,7 +93,7 @@ export async function POST(req: Request) {
const result = streamText({
model: llmgateway.chat(selectedModel),
messages: convertToModelMessages(messages),
// tools,
tools,
});

return result.toUIMessageStreamResponse({
Expand Down
80 changes: 80 additions & 0 deletions apps/playground/src/components/ai-elements/prompt-input.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import {
SquareIcon,
XIcon,
} from "lucide-react";
import { Github } from "lucide-react";
import { nanoid } from "nanoid";
import {
type ChangeEventHandler,
Expand All @@ -31,12 +32,20 @@ import {
} from "react";

import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
Expand Down Expand Up @@ -532,6 +541,77 @@ export const PromptInputTools = ({
/>
);

// Inline Connectors button + dialog (GitHub for now)
export const PromptInputConnectorsButton = ({
className,
}: {
className?: string;
}) => {
const [open, setOpen] = useState(false);
const [token, setToken] = useState<string>(
typeof window !== "undefined"
? (localStorage.getItem("llmgateway_github_token") ?? "")
: "",
);

const save = () => {
if (typeof window !== "undefined") {
if (token) {
localStorage.setItem("llmgateway_github_token", token);
} else {
localStorage.removeItem("llmgateway_github_token");
}
}
setOpen(false);
};
Comment on lines +557 to +566
Copy link
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

Dispatch the GitHub token change event after persisting

useApiKey relies on the llmgateway_github_token_changed custom event to refresh githubToken. Saving here only mutates localStorage, so components using useApiKey keep the stale value, meaning chat requests continue without the new token until the page is reloaded. Please emit the event (and on clear as well) so in-tab consumers update immediately.

 const save = () => {
   if (typeof window !== "undefined") {
     if (token) {
       localStorage.setItem("llmgateway_github_token", token);
     } else {
       localStorage.removeItem("llmgateway_github_token");
     }
+    window.dispatchEvent(new Event("llmgateway_github_token_changed"));
   }
   setOpen(false);
 };
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const save = () => {
if (typeof window !== "undefined") {
if (token) {
localStorage.setItem("llmgateway_github_token", token);
} else {
localStorage.removeItem("llmgateway_github_token");
}
}
setOpen(false);
};
const save = () => {
if (typeof window !== "undefined") {
if (token) {
localStorage.setItem("llmgateway_github_token", token);
} else {
localStorage.removeItem("llmgateway_github_token");
}
window.dispatchEvent(new Event("llmgateway_github_token_changed"));
}
setOpen(false);
};
🤖 Prompt for AI Agents
In apps/playground/src/components/ai-elements/prompt-input.tsx around lines 557
to 566, saving/removing the GitHub token only updates localStorage which leaves
in-tab consumers using useApiKey stale; after calling localStorage.setItem or
localStorage.removeItem dispatch a browser event so listeners refresh: create
and dispatch a CustomEvent named "llmgateway_github_token_changed" (only when
window is defined) immediately after persisting or clearing the token so other
components update their githubToken without a reload, then close the modal as
before.


return (
<>
<PromptInputButton
className={className}
onClick={() => setOpen(true)}
variant="ghost"
>
<Github className="size-4" />
<span className="hidden sm:inline">Connectors</span>
</PromptInputButton>
<Dialog open={open} onOpenChange={setOpen}>
<DialogContent className="sm:max-w-[420px]">
<DialogHeader>
<DialogTitle>Connectors</DialogTitle>
</DialogHeader>
<div className="space-y-4">
<div className="flex items-center gap-2">
<Github className="h-4 w-4" />
<h4 className="font-medium">GitHub</h4>
</div>
<p className="text-sm text-muted-foreground">
Add a GitHub personal access token to enable GitHub MCP tools.
Stored locally.
</p>
<div className="space-y-2">
<Label htmlFor="gh-token">GitHub Token</Label>
<Input
id="gh-token"
type="password"
placeholder="ghp_..."
value={token}
onChange={(e) => setToken(e.target.value)}
/>
<div className="flex justify-end gap-2 pt-2">
<Button variant="outline" onClick={() => setOpen(false)}>
Cancel
</Button>
<Button onClick={save}>Save</Button>
</div>
</div>
</div>
</DialogContent>
</Dialog>
</>
);
};

export type PromptInputButtonProps = ComponentProps<typeof Button>;

export const PromptInputButton = ({
Expand Down
27 changes: 0 additions & 27 deletions apps/playground/src/components/auth/user-provider.tsx

This file was deleted.

15 changes: 13 additions & 2 deletions apps/playground/src/components/playground/chat-ui.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import {
PromptInputTools,
PromptInputToolbar,
PromptInputSubmit,
PromptInputConnectorsButton,
} from "@/components/ai-elements/prompt-input";
import {
Reasoning,
Expand Down Expand Up @@ -171,7 +172,7 @@ export const ChatUI = ({
.join("");
const toolParts = m.parts.filter(
(p) => p.type === "dynamic-tool",
) as any[];
);
const reasoningContent = m.parts
.filter((p) => p.type === "reasoning")
.map((p) => p.text)
Expand Down Expand Up @@ -202,7 +203,11 @@ export const ChatUI = ({

{toolParts.map((tool) => (
<Tool key={tool.toolCallId}>
<ToolHeader type={tool.type} state={tool.state} />
<ToolHeader
type={`tool-${tool.toolName}`}
state={tool.state}
title={tool.toolName}
/>
<ToolContent>
<ToolInput input={tool.input} />
<ToolOutput
Expand Down Expand Up @@ -302,6 +307,11 @@ export const ChatUI = ({
body: {
apiKey: userApiKey,
model: selectedModel,
// Pass GitHub MCP token if available
githubToken:
(typeof window !== "undefined"
? localStorage.getItem("llmgateway_github_token")
: null) || undefined,
},
},
);
Expand Down Expand Up @@ -334,6 +344,7 @@ export const ChatUI = ({
<PromptInputActionAddAttachments />
</PromptInputActionMenuContent>
</PromptInputActionMenu>
<PromptInputConnectorsButton />
</PromptInputTools>
<div className="flex items-center gap-2">
{status === "streaming" ? (
Expand Down
29 changes: 29 additions & 0 deletions apps/playground/src/hooks/useApiKey.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,21 @@ import { toast } from "sonner";

const API_KEY_STORAGE_KEY = "llmgateway_user_api_key";
const API_KEY_CHANGED_EVENT = "llmgateway_api_key_changed";
const GITHUB_TOKEN_STORAGE_KEY = "llmgateway_github_token";
const GITHUB_TOKEN_CHANGED_EVENT = "llmgateway_github_token_changed";

export function useApiKey() {
const [apiKey, setApiKey] = useState<string | null>(null);
const [githubToken, setGithubToken] = useState<string | null>(null);
const [isLoaded, setIsLoaded] = useState(false);

useEffect(() => {
const syncKey = () => {
try {
const storedKey = localStorage.getItem(API_KEY_STORAGE_KEY);
setApiKey(storedKey);
const gh = localStorage.getItem(GITHUB_TOKEN_STORAGE_KEY);
setGithubToken(gh);
} catch {
toast.error("Failed to sync API key from localStorage");
}
Expand All @@ -28,10 +33,12 @@ export function useApiKey() {
window.addEventListener("storage", syncKey);
// Listen for changes from the same tab
window.addEventListener(API_KEY_CHANGED_EVENT, syncKey);
window.addEventListener(GITHUB_TOKEN_CHANGED_EVENT, syncKey);

return () => {
window.removeEventListener("storage", syncKey);
window.removeEventListener(API_KEY_CHANGED_EVENT, syncKey);
window.removeEventListener(GITHUB_TOKEN_CHANGED_EVENT, syncKey);
};
}, []); // Run once on mount

Expand All @@ -54,10 +61,32 @@ export function useApiKey() {
}
}, []);

const setGithubMcpToken = useCallback((token: string) => {
try {
localStorage.setItem(GITHUB_TOKEN_STORAGE_KEY, token);
window.dispatchEvent(new Event(GITHUB_TOKEN_CHANGED_EVENT));
} catch (error) {
toast.error("Failed to save GitHub token to localStorage");
throw error;
}
}, []);

const clearGithubMcpToken = useCallback(() => {
try {
localStorage.removeItem(GITHUB_TOKEN_STORAGE_KEY);
window.dispatchEvent(new Event(GITHUB_TOKEN_CHANGED_EVENT));
} catch {
toast.error("Failed to clear GitHub token from localStorage");
}
}, []);

return {
userApiKey: apiKey,
githubToken,
isLoaded,
setUserApiKey,
clearUserApiKey,
setGithubMcpToken,
clearGithubMcpToken,
};
}
Loading