Refactor CreateApiKeyDialog to use instance IDs

.vscode/launch.json

@@ -14,6 +14,7 @@
                 "GO_ENV": "development",
                 "LLAMACTL_CONFIG_PATH": "${workspaceFolder}/llamactl.dev.yaml"
             },
+            "console": "integratedTerminal",
         }
     ]
 }

cmd/server/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"llamactl/pkg/config"
 	"llamactl/pkg/database"
@@ -11,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 )
 
 // version is set at build time using -ldflags "-X main.version=1.0.0"
@@ -116,14 +118,23 @@ func main() {
 	<-stop
 	fmt.Println("Shutting down server...")
 
-	if err := server.Close(); err != nil {
+	// Create shutdown context with timeout
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer shutdownCancel()
+
+	// Shutdown HTTP server gracefully
+	if err := server.Shutdown(shutdownCtx); err != nil {
 		log.Printf("Error shutting down server: %v\n", err)
 	} else {
 		fmt.Println("Server shut down gracefully.")
 	}
 
-	// Wait for all instances to stop
+	// Stop all instances and cleanup
 	instanceManager.Shutdown()
 
+	if err := db.Close(); err != nil {
+		log.Printf("Error closing database: %v\n", err)
+	}
+
 	fmt.Println("Exiting llamactl.")
 }

pkg/database/database.go

@@ -107,6 +107,12 @@ func Open(config *Config) (*sqliteDB, error) {
 func (db *sqliteDB) Close() error {
 	if db.DB != nil {
 		log.Println("Closing database connection")
+
+		// Checkpoint WAL to merge changes back to main database file
+		if _, err := db.DB.Exec("PRAGMA wal_checkpoint(TRUNCATE)"); err != nil {
+			log.Printf("Warning: Failed to checkpoint WAL: %v", err)
+		}
+
 		return db.DB.Close()
 	}
 	return nil

pkg/database/instances.go

@@ -45,7 +45,7 @@ func (db *sqliteDB) Create(ctx context.Context, inst *instance.Instance) error {
 		) VALUES (?, ?, ?, ?, ?, ?)
 	`
 
-	_, err = db.DB.ExecContext(ctx, query,
+	result, err := db.DB.ExecContext(ctx, query,
 		row.Name, row.Status, row.CreatedAt, row.UpdatedAt, row.OptionsJSON, row.OwnerUserID,
 	)
 
@@ -53,6 +53,14 @@ func (db *sqliteDB) Create(ctx context.Context, inst *instance.Instance) error {
 		return fmt.Errorf("failed to insert instance: %w", err)
 	}
 
+	// Get the auto-generated ID and set it on the instance
+	id, err := result.LastInsertId()
+	if err != nil {
+		return fmt.Errorf("failed to get last insert ID: %w", err)
+	}
+
+	inst.ID = int(id)
+
 	return nil
 }
 
@@ -263,6 +271,7 @@ func (db *sqliteDB) rowToInstance(row *instanceRow) (*instance.Instance, error)
 
 	// Build complete instance JSON with all fields
 	instanceJSON, err := json.Marshal(map[string]any{
+		"id":      row.ID,
 		"name":    row.Name,
 		"created": row.CreatedAt,
 		"status":  row.Status,

pkg/manager/manager.go

@@ -114,11 +114,6 @@ func (im *instanceManager) Shutdown() {
 		}
 		wg.Wait()
 		fmt.Println("All instances stopped.")
-
-		// 4. Close database connection
-		if err := im.db.Close(); err != nil {
-			log.Printf("Error closing database: %v\n", err)
-		}
 	})
 }
 
@@ -181,6 +176,7 @@ func (im *instanceManager) loadInstance(persistedInst *instance.Instance) error
 	inst := instance.New(name, im.globalConfig, options, statusCallback)
 
 	// Restore persisted fields that NewInstance doesn't set
+	inst.ID = persistedInst.ID
 	inst.Created = persistedInst.Created
 	inst.SetStatus(persistedInst.GetStatus())
 

pkg/manager/operations.go

@@ -37,7 +37,6 @@ func (im *instanceManager) ListInstances() ([]*instance.Instance, error) {
 		if node := im.getNodeForInstance(inst); node != nil {
 			remoteInst, err := im.remote.getInstance(ctx, node, inst.Name)
 			if err != nil {
-				// Log error but continue with stale data
 				// Don't fail the entire list operation due to one remote failure
 				continue
 			}

pkg/server/handlers_auth.go

@@ -11,17 +11,12 @@ import (
 	"github.com/go-chi/chi/v5"
 )
 
-// InstancePermission defines the permissions for an API key on a specific instance.
-type InstancePermission struct {
-	InstanceID int `json:"instance_id"`
-}
-
 // CreateKeyRequest represents the request body for creating a new API key.
 type CreateKeyRequest struct {
-	Name                string
-	PermissionMode      auth.PermissionMode
-	ExpiresAt           *int64
-	InstancePermissions []InstancePermission
+	Name           string              `json:"name"`
+	PermissionMode auth.PermissionMode `json:"permission_mode"`
+	ExpiresAt      *int64              `json:"expires_at,omitempty"`
+	InstanceIDs    []int               `json:"instance_ids,omitempty"`
 }
 
 // CreateKeyResponse represents the response returned when creating a new API key.
@@ -87,8 +82,8 @@ func (h *Handler) CreateKey() http.HandlerFunc {
 			writeError(w, http.StatusBadRequest, "invalid_permission_mode", "Permission mode must be 'allow_all' or 'per_instance'")
 			return
 		}
-		if req.PermissionMode == auth.PermissionModePerInstance && len(req.InstancePermissions) == 0 {
-			writeError(w, http.StatusBadRequest, "missing_permissions", "Instance permissions required when permission mode is 'per_instance'")
+		if req.PermissionMode == auth.PermissionModePerInstance && len(req.InstanceIDs) == 0 {
+			writeError(w, http.StatusBadRequest, "missing_permissions", "Instance IDs required when permission mode is 'per_instance'")
 			return
 		}
 		if req.ExpiresAt != nil && *req.ExpiresAt <= time.Now().Unix() {
@@ -108,9 +103,9 @@ func (h *Handler) CreateKey() http.HandlerFunc {
 				instanceIDMap[inst.ID] = true
 			}
 
-			for _, perm := range req.InstancePermissions {
-				if !instanceIDMap[perm.InstanceID] {
-					writeError(w, http.StatusBadRequest, "invalid_instance_id", fmt.Sprintf("Instance ID %d does not exist", perm.InstanceID))
+			for _, instanceID := range req.InstanceIDs {
+				if !instanceIDMap[instanceID] {
+					writeError(w, http.StatusBadRequest, "invalid_instance_id", fmt.Sprintf("Instance ID %d does not exist", instanceID))
 					return
 				}
 			}
@@ -142,12 +137,12 @@ func (h *Handler) CreateKey() http.HandlerFunc {
 			UpdatedAt:      now,
 		}
 
-		// Convert InstancePermissions to KeyPermissions
+		// Convert InstanceIDs to KeyPermissions
 		var keyPermissions []auth.KeyPermission
-		for _, perm := range req.InstancePermissions {
+		for _, instanceID := range req.InstanceIDs {
 			keyPermissions = append(keyPermissions, auth.KeyPermission{
 				KeyID:      0, // Will be set by database after key creation
-				InstanceID: perm.InstanceID,
+				InstanceID: instanceID,
 			})
 		}
 

webui/src/components/apikeys/CreateApiKeyDialog.tsx

@@ -8,9 +8,9 @@ import { Checkbox } from "@/components/ui/checkbox";
 import { Alert, AlertDescription } from "@/components/ui/alert";
 import { Loader2 } from "lucide-react";
 import { apiKeysApi } from "@/lib/api";
-import { CreateKeyRequest, PermissionMode, InstancePermission } from "@/types/apiKey";
+import { PermissionMode, type CreateKeyRequest } from "@/types/apiKey";
 import { useInstances } from "@/contexts/InstancesContext";
-import { format, addDays } from "date-fns";
+import { format } from "date-fns";
 
 interface CreateApiKeyDialogProps {
   open: boolean;
@@ -61,22 +61,19 @@ function CreateApiKeyDialog({ open, onOpenChange, onKeyCreated }: CreateApiKeyDi
     }
 
     // Build request
-    const permissions: InstancePermission[] = [];
+    const instanceIds: number[] = [];
     if (permissionMode === PermissionMode.PerInstance) {
-      Object.entries(instancePermissions).forEach(([instanceId, canInfer]) => {
-        if (canInfer) {
-          permissions.push({
-            InstanceID: parseInt(instanceId),
-            CanInfer: true,
-          });
+      Object.entries(instancePermissions).forEach(([instanceId, hasPermission]) => {
+        if (hasPermission) {
+          instanceIds.push(parseInt(instanceId));
         }
       });
     }
 
     const request: CreateKeyRequest = {
-      Name: name.trim(),
-      PermissionMode: permissionMode,
-      InstancePermissions: permissions,
+      name: name.trim(),
+      permission_mode: permissionMode,
+      instance_ids: instanceIds,
     };
 
     // Add expiration if provided
@@ -87,7 +84,7 @@ function CreateApiKeyDialog({ open, onOpenChange, onKeyCreated }: CreateApiKeyDi
         setError("Expiration date must be in the future");
         return;
       }
-      request.ExpiresAt = Math.floor(expirationDate.getTime() / 1000);
+      request.expires_at = Math.floor(expirationDate.getTime() / 1000);
     }
 
     setLoading(true);
@@ -107,10 +104,10 @@ function CreateApiKeyDialog({ open, onOpenChange, onKeyCreated }: CreateApiKeyDi
   };
 
   const handleInstancePermissionChange = (instanceId: number, checked: boolean) => {
-    setInstancePermissions({
-      ...instancePermissions,
+    setInstancePermissions(prev => ({
+      ...prev,
       [instanceId]: checked,
-    });
+    }));
   };
 
   return (
@@ -172,14 +169,19 @@ function CreateApiKeyDialog({ open, onOpenChange, onKeyCreated }: CreateApiKeyDi
                   <p className="text-sm text-muted-foreground">No instances available</p>
                 ) : (
                   <div className="space-y-2">
-                    {instances.map((instance) => (
-                      <div key={instance.id} className="flex items-center space-x-2">
+                    {instances.map((instance, index) => {
+                      const isChecked = !!instancePermissions[instance.id];
+                      return (
+                        <div
+                          key={`${instance.name}-${index}`}
+                          className="flex items-center space-x-2"
+                        >
                           <Checkbox
                             id={`instance-${instance.id}`}
-                          checked={instancePermissions[instance.id] || false}
-                          onCheckedChange={(checked) =>
-                            handleInstancePermissionChange(instance.id, checked as boolean)
-                          }
+                            checked={isChecked}
+                            onCheckedChange={(checked) => {
+                              handleInstancePermissionChange(instance.id, checked as boolean);
+                            }}
                             disabled={loading}
                           />
                           <Label
@@ -188,9 +190,9 @@ function CreateApiKeyDialog({ open, onOpenChange, onKeyCreated }: CreateApiKeyDi
                           >
                             {instance.name}
                           </Label>
-                        <span className="text-sm text-muted-foreground">Can Infer</span>
                         </div>
-                    ))}
+                      );
+                    })}
                   </div>
                 )}
               </div>

webui/src/components/settings/ApiKeysSection.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useState } from "react";
+import { useEffect, useState, Fragment } from "react";
 import { Button } from "@/components/ui/button";
 import { Badge } from "@/components/ui/badge";
 import { Alert, AlertDescription } from "@/components/ui/alert";
@@ -175,9 +175,8 @@ function ApiKeysSection() {
             </thead>
             <tbody>
               {keys.map((key) => (
-                <>
+                <Fragment key={key.id}>
                   <tr
-                    key={key.id}
                     className="border-t hover:bg-muted/50 cursor-pointer"
                     onClick={() => handleRowClick(key)}
                   >
@@ -236,25 +235,15 @@ function ApiKeysSection() {
                           <p className="text-sm text-muted-foreground">Loading permissions...</p>
                         ) : permissions[key.id] ? (
                           <div className="space-y-2">
-                            <p className="text-sm font-semibold">Instance Permissions:</p>
-                            <table className="w-full text-sm">
-                              <thead>
-                                <tr className="border-b">
-                                  <th className="text-left py-2">Instance Name</th>
-                                  <th className="text-left py-2">Can Infer</th>
-                                </tr>
-                              </thead>
-                              <tbody>
+                            <p className="text-sm font-semibold">Allowed Instances:</p>
+                            <ul className="text-sm space-y-1">
                               {permissions[key.id].map((perm) => (
-                                  <tr key={perm.instance_id} className="border-b">
-                                    <td className="py-2">{perm.instance_name}</td>
-                                    <td className="py-2">
-                                      <Check className="h-4 w-4 text-green-600" />
-                                    </td>
-                                  </tr>
+                                <li key={perm.instance_id} className="flex items-center gap-2">
+                                  <Check className="h-3 w-3 text-green-600" />
+                                  {perm.instance_name}
+                                </li>
                               ))}
-                              </tbody>
-                            </table>
+                            </ul>
                           </div>
                         ) : (
                           <p className="text-sm text-muted-foreground">No permissions data</p>
@@ -262,7 +251,7 @@ function ApiKeysSection() {
                       </td>
                     </tr>
                   )}
-                </>
+                </Fragment>
               ))}
             </tbody>
           </table>

webui/src/types/apiKey.ts

@@ -15,14 +15,10 @@ export interface ApiKey {
 }
 
 export interface CreateKeyRequest {
-  Name: string
-  PermissionMode: PermissionMode
-  ExpiresAt?: number
-  InstancePermissions: InstancePermission[]
-}
-
-export interface InstancePermission {
-  InstanceID: number
+  name: string
+  permission_mode: PermissionMode
+  expires_at?: number
+  instance_ids: number[]
 }
 
 export interface CreateKeyResponse extends ApiKey {