Update app.py

app.py

@@ -28,13 +28,14 @@ CREATOR_APP_PY_URL = f"{CREATOR_URL}/blob/main/app.py"
 SERVER_APP_PY_URL = f"{SERVER_URL}/blob/main/app.py"
 
 # ==============================================================================
-#  API CALL WRAPPER FUNCTIONS (Your working versions)
+#  API CALL WRAPPER FUNCTIONS (CORRECTED)
 # ==============================================================================
 
 def get_creator_endpoints():
     """Fetches the list of supported endpoints by making an HTTP request to the Creator's JSON file."""
     status = f"Fetching endpoint list from {CREATOR_ENDPOINTS_JSON_URL}..."
-    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] # Initial state
+    # Using yield for streaming status updates
+    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] 
     try:
         response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
         response.raise_for_status()
@@ -53,56 +54,66 @@ def create_image_via_api(service_name: str, secret_data: str, available_endpoint
         raise gr.Error("Please select a service and provide secret data.")
     
     status = f"Looking up public key for '{service_name}'..."
-    yield None, None, status
+    # No yield here, we'll return everything at the end.
 
-    public_key = next((e['public_key'] for e in available_endpoints if e['name'] == service_name), None)
-    if not public_key:
-        raise gr.Error(f"Could not find public key for '{service_name}' in the fetched configuration.")
-        
-    status = f"Connecting to Creator: {CREATOR_SPACE_ID}..."
-    yield None, None, status
-    
     try:
+        public_key = next((e['public_key'] for e in available_endpoints if e['name'] == service_name), None)
+        if not public_key:
+            raise gr.Error(f"Could not find public key for '{service_name}' in the fetched configuration.")
+            
+        status = f"Connecting to Creator: {CREATOR_SPACE_ID}..."
+        logger.info(status)
+        
         client = Client(src=CREATOR_SPACE_ID)
         temp_filepath = client.predict(secret_data, public_key, api_name="/create_image")
         
-        if not temp_filepath: raise gr.Error("Creator API did not return an image.")
+        if not temp_filepath: 
+            raise gr.Error("Creator API did not return an image.")
         
         # --- PNG FIX ---
-        # The API returns a path to a temp file which could be any format.
-        # We open it, ensure it's saved as PNG, and return that path.
+        # Load the image from the temp path and return the PIL Image object directly.
+        # The gr.Image component's `format="png"` will handle the rest.
         created_image = Image.open(temp_filepath)
         
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as png_file:
-            created_image.save(png_file.name, "PNG")
-            png_filepath = png_file.name
-            
         status = f"✅ Success! Image created for '{service_name}'."
-        yield created_image, png_filepath, status
+        return created_image, status
 
     except Exception as e:
         logger.error(f"Creator API call failed: {e}", exc_info=True)
-        yield None, None, f"❌ Error calling Creator API: {e}"
+        # On error, return an empty image and the error status.
+        return None, f"❌ Error calling Creator API: {e}"
 
 def decrypt_image_via_api(image: Image.Image):
     """Calls the Server Space API to decrypt an image."""
-    if image is None: raise gr.Error("Please upload an image to decrypt.")
+    if image is None: 
+        raise gr.Error("Please upload an image to decrypt.")
+        
     status = f"Connecting to Server: {SERVER_SPACE_ID}..."
-    yield None, status
+    logger.info(status)
+    
     try:
         client = Client(src=SERVER_SPACE_ID)
         with io.BytesIO() as buffer:
+            # Explicitly save as PNG to ensure format is correct before base64 encoding
             image.save(buffer, format="PNG")
             b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
         
         status = f"Calling API on {SERVER_SPACE_ID}..."
-        yield None, status
+        logger.info(status)
+        
         decrypted_json = client.predict(b64_string, api_name="/keylock-auth-decoder")
+        
         status = "✅ Success! Data decrypted by the Server."
-        yield decrypted_json, status
+        logger.info(f"Decryption successful. Data: {decrypted_json}")
+        
+        # --- DECRYPTED DATA FIX ---
+        # Use a single, final return statement. This is more robust.
+        return decrypted_json, status
+
     except Exception as e:
         logger.error(f"Server API call failed: {e}", exc_info=True)
-        yield None, f"❌ Error calling Server API: {e}"
+        # Return an empty dict and the error message for a clean UI update.
+        return {}, f"❌ Error calling Server API: {e}"
 
 def generate_rsa_keys():
     """Generates a new RSA key pair."""
@@ -117,7 +128,7 @@ def generate_rsa_keys():
     return private_pem, public_pem
 
 # ==============================================================================
-#  GRADIO DASHBOARD INTERFACE
+#  GRADIO DASHBOARD INTERFACE (Corrected Layout and Components)
 # ==============================================================================
 theme = gr.themes.Base(
     primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
@@ -142,15 +153,14 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                 with gr.Column(scale=2):
                     with gr.Row():
                         creator_service_dropdown = gr.Dropdown(label="Target Service", interactive=True, info="Select the API server you want to encrypt data for.")
-                        refresh_button = gr.Button("🔄", scale=0, size="sm")
+                        refresh_button = gr.Button("🔄", scale=0, size="sm", tooltip="Refresh Target Service List")
                     creator_secret_input = gr.Textbox(lines=8, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
                     creator_button = gr.Button("✨ Create Auth Image via API", variant="primary")
                 with gr.Column(scale=1):
                     creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=False) # Download handled by gr.File
                     # --- PNG FIX ---
-                    # Use a dedicated, visible File component for reliable PNG downloads.
-                    creator_download_output = gr.File(label="Download Image as PNG", interactive=False)
+                    # The format='png' argument tells the component's download button to create a PNG.
+                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True, format="png")
 
         with gr.TabItem("② Client / Decoder", id=1):
             gr.Markdown("## Decrypt an Authentication Image")
@@ -186,6 +196,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
     def refresh_endpoints():
+        # This is a generator function, so we need to iterate to get the last value.
         *_, last_yield = get_creator_endpoints()
         return last_yield
 
@@ -195,9 +206,13 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     creator_button.click(
         fn=create_image_via_api,
         inputs=[creator_service_dropdown, creator_secret_input, endpoints_state],
-        outputs=[creator_image_output, creator_download_output, creator_status] # Now outputs to the gr.File component
+        outputs=[creator_image_output, creator_status] # Removed the gr.File output
+    )
+    client_button.click(
+        fn=decrypt_image_via_api,
+        inputs=[client_image_input],
+        outputs=[client_json_output, client_status]
     )
-    client_button.click(fn=decrypt_image_via_api, inputs=[client_image_input], outputs=[client_json_output, client_status])
 
 if __name__ == "__main__":
     demo.launch()