Fix headless lifecycle and array type setting

README.md

@@ -142,6 +142,15 @@ For a binary that is already imported into the project, use `-process` instead:
   -preScript GAMCPStartServerScript.java "host=127.0.0.1" "port=8080"
 ```
 
+To keep a headless MCP session open after analysis completes, run the server as a post-script with wait mode:
+
+```bash
+"$GHIDRA_INSTALL_DIR/support/analyzeHeadless" /tmp/ghidra-projects McpHeadless \
+  -process binary_name \
+  -scriptPath "$GHIDRASSISTMCP_EXT/ghidra_scripts" \
+  -postScript GAMCPStartServerScript.java "host=127.0.0.1" "port=8080" "wait=true"
+```
+
 MCP clients can connect to:
 
 ```text
@@ -150,7 +159,7 @@ SSE messages:    http://127.0.0.1:8080/message
 Streamable HTTP: http://127.0.0.1:8080/mcp
 ```
 
-The headless MCP server runs inside the `analyzeHeadless` JVM and uses the loaded `currentProgram`. Keep that process alive while clients are connected; when `analyzeHeadless` exits, the MCP server exits with it.
+The headless MCP server runs inside the `analyzeHeadless` JVM and uses the loaded `currentProgram`. The server holds a program consumer while it is running so MCP requests do not race against program database closure. Use `wait=true` when you want `analyzeHeadless` to stay open for interactive MCP clients; cancel the script or terminate the process to stop the server.
 
 ## Available Tools
 
@@ -179,6 +188,8 @@ GhidrAssistMCP provides 38 tools organized into categories. Several tools use an
 | `get_current_function` | Get function at current cursor position |
 | `get_function_stack_layout` | Get stack frame layout with variable offsets |
 | `get_basic_blocks` | Get basic block information for a function |
+| `create_function` | Create/define a function at an address |
+| `disassemble_at` | Disassemble code at an address |
 
 ### Binary Information
 
@@ -275,7 +286,7 @@ Rename multiple symbols in one operation.
 | ------ | ----------- |
 | `list` | List all available data types |
 | `get_info` | Get detailed data type information and structure definitions |
-| `set` | Set data type at a specific address |
+| `set` | Set data type at a specific address, including arrays with `array_count` or suffix syntax like `int[16]` |
 | `delete` | Delete a data type by name (optionally scoped by `category`) |
 
 #### `bookmarks` - Bookmark Management Tool
@@ -454,6 +465,54 @@ If multiple types share the same name across categories, pass `category` (or pas
 }
 ```
 
+### Set an Array Data Type
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "types",
+    "arguments": {
+      "action": "set",
+      "address": "0x00402000",
+      "data_type": "int[16]"
+    }
+  }
+}
+```
+
+Equivalent form:
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "types",
+    "arguments": {
+      "action": "set",
+      "address": "0x00402000",
+      "data_type": "int",
+      "array_count": 16
+    }
+  }
+}
+```
+
+### Create a Function
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "create_function",
+    "arguments": {
+      "address": "0x00401000",
+      "name": "mainWndProc"
+    }
+  }
+}
+```
+
 ### Rename Function (Action-Based)
 
 ```json
@@ -592,7 +651,7 @@ GhidrAssistMCP/
 - `rename_symbol`: `target_type: function|data|variable`
 - `comments`: `action: get|set|list|remove`
 - `variables`: `action: list|rename|set_type|set_prototype` with `scope: auto|local|global` for rename
-- `types`: `action: list|get_info|set|delete`
+- `types`: `action: list|get|set|create_struct|create_enum|create_typedef|delete`
 - `bookmarks`: `action: list|set|remove`
 - `xrefs`: `address|function` with `include_calls` parameter
 

src/main/java/ghidrassistmcp/GhidrAssistMCPHeadlessServer.java

@@ -81,6 +81,10 @@ public class GhidrAssistMCPHeadlessServer {
             Msg.info(this, "Headless MCP server stopped");
         } catch (Exception e) {
             Msg.error(this, "Error stopping headless MCP server", e);
+        } finally {
+            if (headlessBackend != null) {
+                headlessBackend.shutdown();
+            }
         }
         server = null;
         headlessBackend = null;
@@ -111,35 +115,64 @@ public class GhidrAssistMCPHeadlessServer {
     private static class HeadlessBackend extends GhidrAssistMCPBackend {
 
         private volatile Program currentProgram;
+        private final Object programConsumer = new Object();
 
         HeadlessBackend(Program program) {
             super();
-            this.currentProgram = program;
-            // Notify backend of the initial program
-            onProgramActivated(program);
+            setProgram(program);
         }
 
-        void setProgram(Program program) {
-            if (this.currentProgram != null) {
-                onProgramDeactivated(this.currentProgram);
+        synchronized void setProgram(Program program) {
+            Program oldProgram = this.currentProgram;
+            if (oldProgram != null) {
+                onProgramDeactivated(oldProgram);
+                releaseProgram(oldProgram);
             }
             this.currentProgram = program;
             if (program != null) {
+                if (!program.isClosed()) {
+                    program.addConsumer(programConsumer);
+                }
                 onProgramActivated(program);
             }
         }
 
         @Override
         public Program getCurrentProgram() {
+            if (currentProgram != null && currentProgram.isClosed()) {
+                Msg.warn(this, "Headless current program is closed; clearing program reference");
+                currentProgram = null;
+            }
             return currentProgram;
         }
 
         @Override
         public List<Program> getAllOpenPrograms() {
-            if (currentProgram != null) {
-                return Collections.singletonList(currentProgram);
+            Program program = getCurrentProgram();
+            if (program != null) {
+                return Collections.singletonList(program);
             }
             return Collections.emptyList();
         }
+
+        synchronized void shutdown() {
+            Program program = currentProgram;
+            currentProgram = null;
+            if (program != null) {
+                onProgramDeactivated(program);
+                releaseProgram(program);
+            }
+            getTaskManager().shutdown();
+        }
+
+        private void releaseProgram(Program program) {
+            try {
+                if (!program.isClosed() && program.isUsedBy(programConsumer)) {
+                    program.release(programConsumer);
+                }
+            } catch (Exception e) {
+                Msg.warn(this, "Failed to release headless program consumer: " + e.getMessage(), e);
+            }
+        }
     }
 }

src/main/java/ghidrassistmcp/scripts/GAMCPStartServerScript.java

@@ -24,8 +24,9 @@ public class GAMCPStartServerScript extends GhidraScript {
 
         String host = "localhost";
         int port = 8080;
+        boolean waitForClients = false;
 
-        // Parse optional arguments: host=... port=...
+        // Parse optional arguments: host=... port=... wait=true|false
         String[] args = getScriptArgs();
         if (args != null) {
             for (String arg : args) {
@@ -37,6 +38,8 @@ public class GAMCPStartServerScript extends GhidraScript {
                     } catch (NumberFormatException e) {
                         Msg.warn(this, "Invalid port argument, using default 8080");
                     }
+                } else if (arg.startsWith("wait=")) {
+                    waitForClients = Boolean.parseBoolean(arg.substring(5));
                 }
             }
         }
@@ -46,11 +49,30 @@ public class GAMCPStartServerScript extends GhidraScript {
         if (mcpServer.isRunning()) {
             Msg.info(this, "MCP server already running, updating program reference");
             mcpServer.setProgram(currentProgram);
+            if (waitForClients) {
+                waitUntilCancelled(mcpServer);
+            }
             return;
         }
 
         Msg.info(this, "Starting headless MCP server for: " + currentProgram.getName());
         mcpServer.start(currentProgram, host, port);
         Msg.info(this, "Headless MCP server ready on " + host + ":" + port);
+        if (waitForClients) {
+            waitUntilCancelled(mcpServer);
+        }
+    }
+
+    private void waitUntilCancelled(GhidrAssistMCPHeadlessServer mcpServer) {
+        Msg.info(this, "Headless MCP server wait mode enabled; cancel the script or terminate analyzeHeadless to stop");
+        try {
+            while (!monitor.isCancelled() && mcpServer.isRunning()) {
+                Thread.sleep(1000);
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        } finally {
+            mcpServer.stop();
+        }
     }
 }

src/main/java/ghidrassistmcp/tools/DataTypeResolver.java

@@ -0,0 +1,125 @@
+package ghidrassistmcp.tools;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import ghidra.program.model.data.ArrayDataType;
+import ghidra.program.model.data.DataType;
+import ghidra.program.model.data.DataTypeManager;
+
+final class DataTypeResolver {
+
+    private static final Pattern ARRAY_SUFFIX = Pattern.compile("^(.+)\\[(\\d+)]\\s*$");
+
+    private DataTypeResolver() {
+    }
+
+    static Result resolve(DataTypeManager dtm, String dataTypeName, Object arrayCountValue) {
+        if (dataTypeName == null || dataTypeName.isBlank()) {
+            return Result.error("Data type name is required");
+        }
+
+        String baseTypeName = dataTypeName.trim();
+        Integer arrayCount = parseArrayCount(arrayCountValue);
+        if (arrayCountValue != null && arrayCount == null) {
+            return Result.error("array_count must be an integer");
+        }
+
+        Matcher matcher = ARRAY_SUFFIX.matcher(baseTypeName);
+        if (matcher.matches()) {
+            if (arrayCount != null) {
+                return Result.error("Specify array size either with data_type suffix syntax or array_count, not both");
+            }
+            baseTypeName = matcher.group(1).trim();
+            arrayCount = Integer.parseInt(matcher.group(2));
+        }
+
+        if (arrayCount != null && arrayCount < 1) {
+            return Result.error("array_count must be >= 1");
+        }
+
+        DataType baseType = resolveBaseType(dtm, baseTypeName);
+        if (baseType == null) {
+            return Result.error("Data type not found: " + baseTypeName);
+        }
+
+        if (arrayCount == null) {
+            return Result.ok(baseType);
+        }
+
+        int elementLength = baseType.getLength();
+        if (elementLength <= 0) {
+            return Result.error("Cannot build array of variable-length type '" + baseType.getName() + "'");
+        }
+
+        return Result.ok(new ArrayDataType(baseType, arrayCount, elementLength));
+    }
+
+    private static Integer parseArrayCount(Object value) {
+        if (value == null) {
+            return null;
+        }
+        if (value instanceof Number) {
+            double doubleValue = ((Number) value).doubleValue();
+            int intValue = ((Number) value).intValue();
+            return doubleValue == intValue ? intValue : null;
+        }
+        if (value instanceof String) {
+            try {
+                return Integer.parseInt(((String) value).trim());
+            } catch (NumberFormatException e) {
+                return null;
+            }
+        }
+        return null;
+    }
+
+    private static DataType resolveBaseType(DataTypeManager dtm, String dataTypeName) {
+        DataType dataType = null;
+        if (dataTypeName.startsWith("/")) {
+            dataType = dtm.getDataType(dataTypeName);
+        }
+        if (dataType == null) {
+            dataType = dtm.getDataType("/" + dataTypeName);
+        }
+        if (dataType == null) {
+            dataType = dtm.getDataType(dataTypeName);
+        }
+        if (dataType != null) {
+            return dataType;
+        }
+
+        List<DataType> allTypes = new ArrayList<>();
+        dtm.getAllDataTypes(allTypes);
+        for (DataType candidate : allTypes) {
+            if (candidate.getName().equals(dataTypeName)) {
+                return candidate;
+            }
+        }
+        return null;
+    }
+
+    static final class Result {
+        final DataType dataType;
+        final String errorMessage;
+
+        private Result(DataType dataType, String errorMessage) {
+            this.dataType = dataType;
+            this.errorMessage = errorMessage;
+        }
+
+        static Result ok(DataType dataType) {
+            return new Result(dataType, null);
+        }
+
+        static Result error(String errorMessage) {
+            return new Result(null, errorMessage);
+        }
+
+        boolean isError() {
+            return errorMessage != null;
+        }
+    }
+}

src/main/java/ghidrassistmcp/tools/SetDataTypeTool.java

@@ -8,7 +8,6 @@ import java.util.Map;
 
 import ghidra.program.model.address.Address;
 import ghidra.program.model.data.DataType;
-import ghidra.program.model.data.DataTypeManager;
 import ghidra.program.model.listing.Program;
 import ghidra.program.model.util.CodeUnitInsertionException;
 import ghidrassistmcp.McpTool;
@@ -44,7 +43,11 @@ public class SetDataTypeTool implements McpTool {
         return new McpSchema.JsonSchema("object", 
             Map.of(
                 "address", new McpSchema.JsonSchema("string", null, null, null, null, null),
-                "data_type", new McpSchema.JsonSchema("string", null, null, null, null, null)
+                "data_type", new McpSchema.JsonSchema("string", null, null, null, null, null),
+                "array_count", Map.of(
+                    "type", "integer",
+                    "description", "Optional: wraps data_type in an array of this many elements. Alternatively use suffix syntax like 'int[16]'."
+                )
             ),
             List.of("address", "data_type"), null, null, null);
     }
@@ -76,32 +79,14 @@ public class SetDataTypeTool implements McpTool {
                 .build();
         }
         
-        // Find the data type
-        DataTypeManager dtm = currentProgram.getDataTypeManager();
-        DataType dataType = dtm.getDataType("/" + dataTypeName);
-        if (dataType == null) {
-            // Try finding by name without path
-            dataType = dtm.getDataType(dataTypeName);
-        }
-        
-        if (dataType == null) {
-             // Search for type in all categories if simple lookup failed
-             var allTypes = new java.util.ArrayList<DataType>();
-             dtm.getAllDataTypes(allTypes);
-             for (DataType dt : allTypes) {
-                 if (dt.getName().equals(dataTypeName)) {
-                     dataType = dt;
-                     break;
-                 }
-             }
-        }
-
-        if (dataType == null) {
+        DataTypeResolver.Result resolvedType =
+            DataTypeResolver.resolve(currentProgram.getDataTypeManager(), dataTypeName, arguments.get("array_count"));
+        if (resolvedType.isError()) {
             return McpSchema.CallToolResult.builder()
-                .addTextContent("Data type not found: " + dataTypeName + 
-                              ". Use built-in types like 'int', 'char', 'void*', etc.")
+                .addTextContent(resolvedType.errorMessage)
                 .build();
         }
+        DataType dataType = resolvedType.dataType;
         
         // Start transaction
         int transactionID = currentProgram.startTransaction("Set Data Type");

src/main/java/ghidrassistmcp/tools/TypesTool.java

@@ -58,7 +58,8 @@ public class TypesTool implements McpTool {
                 Map.entry("offset", Map.of("type", "integer", "description", "Pagination offset (for list)")),
                 Map.entry("limit", Map.of("type", "integer", "description", "Pagination limit (for list, default 100)")),
                 Map.entry("address", Map.of("type", "string", "description", "Address (for set action)")),
-                Map.entry("data_type", Map.of("type", "string", "description", "Data type name (for set action)")),
+                Map.entry("data_type", Map.of("type", "string", "description", "Data type name (for set action). Supports array suffix syntax like 'int[16]'.")),
+                Map.entry("array_count", Map.of("type", "integer", "description", "For set action: wraps data_type in an array of this many elements")),
                 Map.entry("size", Map.of("type", "integer", "description", "Size in bytes (for create_struct/create_enum)")),
                 Map.entry("packed", Map.of("type", "boolean", "description", "Enable packing (for create_struct)")),
                 Map.entry("values", Map.of("type", "object", "description", "Enum values as name:value pairs (for create_enum)")),
@@ -192,10 +193,10 @@ public class TypesTool implements McpTool {
         try { address = program.getAddressFactory().getAddress(addressStr); }
         catch (Exception e) { return result("Invalid address: " + addressStr); }
 
-        DataTypeManager dtm = program.getDataTypeManager();
-        DataType dataType = dtm.getDataType("/" + dataTypeName);
-        if (dataType == null) dataType = dtm.getDataType(dataTypeName);
-        if (dataType == null) return result("Data type not found: " + dataTypeName);
+        DataTypeResolver.Result resolvedType =
+            DataTypeResolver.resolve(program.getDataTypeManager(), dataTypeName, arguments.get("array_count"));
+        if (resolvedType.isError()) return result(resolvedType.errorMessage);
+        DataType dataType = resolvedType.dataType;
 
         int txId = program.startTransaction("Set Data Type");
         try {