proxy traffic perf updates and ui parameter dialog modernization

2026-06-08 14:55:38 +00:00 · 2026-05-08 21:41:09 -05:00
parent f3c5978ee8
commit 98a9925be0
12 changed files with 2332 additions and 890 deletions
@@ -22,6 +22,7 @@ import {getSkewedNow} from "../../utilities/Time";
 import { useTheme } from '@mui/material/styles';
 import {MythicStyledTooltip} from "../../MythicComponents/MythicStyledTooltip";
 import {getReadableTextColor, isValidHexColor} from "../../MythicComponents/MythicColorInput";
+import {copyStringToClipboard} from "../../utilities/Clipboard";

 const GetLoadedCommandsSubscription = gql`
 subscription GetLoadedCommandsSubscription($callback_id: Int!){
@@ -174,22 +175,39 @@ const IsRepeatableCLIParameterType = (parameter_type) => {
 const TaskingContextChip = ({title, label, value, color, callbackColor, emphasized=false}) => {
    const safeColor = isValidHexColor(color) ? color : "#000000";
    const borderColor = isValidHexColor(callbackColor) ? callbackColor : undefined;
-    const tooltipTitle = `${title}: ${value}`;
+    const copyValue = String(value);
+    const copyLabel = `Copy ${title}`;
+    const onCopyValue = (event) => {
+        event.preventDefault();
+        event.stopPropagation();
+        if(copyStringToClipboard(copyValue)){
+            snackActions.success("Copied tasking context");
+        }else{
+            snackActions.error("Failed to copy tasking context");
+        }
+    }
+    const onCopyKeyDown = (event) => {
+        if(event.key === "Enter" || event.key === " "){
+            onCopyValue(event);
+        }
+    }
    return (
-        <MythicStyledTooltip title={tooltipTitle}>
-            <span className={`mythic-tasking-context-chip${emphasized ? " mythic-tasking-context-chip-emphasized" : ""}`}
-                  style={{
-                      backgroundColor: safeColor,
-                      borderColor: borderColor,
-                      color: getReadableTextColor(safeColor),
-                  }}
-                  title={tooltipTitle}>
-                {label !== "" &&
-                    <span className="mythic-tasking-context-chip-label">{label}</span>
-                }
-                <span className="mythic-tasking-context-chip-value">{value}</span>
-            </span>
-        </MythicStyledTooltip>
+        <span className={`mythic-tasking-context-chip${emphasized ? " mythic-tasking-context-chip-emphasized" : ""}`}
+              aria-label={copyLabel}
+              onClick={onCopyValue}
+              onKeyDown={onCopyKeyDown}
+              role="button"
+              style={{
+                  backgroundColor: safeColor,
+                  borderColor: borderColor,
+                  color: getReadableTextColor(safeColor),
+              }}
+              tabIndex={0}>
+            {label !== "" &&
+                <span className="mythic-tasking-context-chip-label">{label}</span>
+            }
+            <span className="mythic-tasking-context-chip-value">{copyValue}</span>
+        </span>
    )
 }

@@ -1,27 +1,23 @@
 import React, {useEffect, useState} from 'react';
-import Button from '@mui/material/Button';
-import DialogActions from '@mui/material/DialogActions';
 import DialogContent from '@mui/material/DialogContent';
 import DialogTitle from '@mui/material/DialogTitle';
 import Typography from '@mui/material/Typography';
-import Table from '@mui/material/Table';
-import TableBody from '@mui/material/TableBody';
-import TableCell from '@mui/material/TableCell';
-import TableContainer from '@mui/material/TableContainer';
-import TableHead from '@mui/material/TableHead';
-import TableRow from '@mui/material/TableRow';
+import Box from '@mui/material/Box';
+import Chip from '@mui/material/Chip';
 import {TaskParametersDialogRow} from './TaskParametersDialogRow';
 import {gql, useLazyQuery, useMutation, useQuery} from '@apollo/client';
 import FormControl from '@mui/material/FormControl';
 import TextField from '@mui/material/TextField';
 import MenuItem from '@mui/material/MenuItem';
-import Input from '@mui/material/Input';
 import {UploadTaskFile} from '../../MythicComponents/MythicFileUpload';
 import {Backdrop, CircularProgress} from '@mui/material';
-import Divider from '@mui/material/Divider';
 import {b64DecodeUnicode} from './ResponseDisplay';
 import {snackActions} from "../../utilities/Snackbar";
-import {MythicDraggableDialogTitle} from "../../MythicComponents/MythicDraggableDialogTitle";
+import TerminalIcon from '@mui/icons-material/Terminal';
+import AdminPanelSettingsIcon from '@mui/icons-material/AdminPanelSettings';
+import SegmentIcon from '@mui/icons-material/Segment';
+import FactCheckIcon from '@mui/icons-material/FactCheck';
+import {MythicDialogButton, MythicDialogFooter} from "../../MythicComponents/MythicDialogLayout";

 //if we need to get all the loaded commands for the callback and filter, use this
 const GetLoadedCommandsQuery = gql`
@@ -1030,27 +1026,69 @@ export function TaskParametersDialog(props) {
        }
        return collapsedParameters;
    }
+  const requiredCount = parameters.filter((param) => param.required).length;
+  const commandName = commandInfo.cmd || props.command?.cmd || "Command";
  return (
    <React.Fragment>
-        <MythicDraggableDialogTitle>{commandInfo.cmd}'s Parameters</MythicDraggableDialogTitle>
-        <DialogContent dividers={true}>
-            <Backdrop open={backdropOpen} style={{zIndex: 2, position: "absolute"}}>
+        <DialogTitle id="mythic-draggable-title" className="mythic-task-parameters-title">
+            <Box className="mythic-task-parameters-title-row">
+                <Box className="mythic-task-parameters-title-icon">
+                    <TerminalIcon fontSize="small" />
+                </Box>
+                <Box className="mythic-task-parameters-title-copy">
+                    <Typography component="div" className="mythic-task-parameters-title-main">
+                        {commandName}
+                    </Typography>
+                    <Typography component="div" className="mythic-task-parameters-title-subtitle">
+                        Command parameters
+                    </Typography>
+                </Box>
+                <Box className="mythic-task-parameters-title-meta">
+                    {commandInfo?.payloadtype?.name &&
+                        <Chip size="small" className="mythic-task-parameters-title-chip" icon={<SegmentIcon />} label={commandInfo.payloadtype.name} />
+                    }
+                    <Chip size="small" className="mythic-task-parameters-title-chip" icon={<FactCheckIcon />} label={`${parameters.length} parameter${parameters.length === 1 ? "" : "s"}`} />
+                    {requiredCount > 0 &&
+                        <Chip size="small" className="mythic-task-parameters-title-chip mythic-task-parameters-title-chip-warning" label={`${requiredCount} required`} />
+                    }
+                    <Chip
+                        size="small"
+                        className={`mythic-task-parameters-title-chip${commandInfo.needs_admin ? " mythic-task-parameters-title-chip-warning" : ""}`}
+                        icon={<AdminPanelSettingsIcon />}
+                        label={commandInfo.needs_admin ? "Admin required" : "No admin"}
+                    />
+                </Box>
+            </Box>
+        </DialogTitle>
+        <DialogContent dividers={true} className="mythic-task-parameters-content">
+            <Backdrop open={backdropOpen} className="mythic-task-parameters-backdrop">
                <CircularProgress color="inherit" />
            </Backdrop>
-            <Typography component="div" >
-                <b>Description</b> <pre style={{margin:0, wordBreak: "break-word", overflow: "word-wrap", whiteSpace: "pre-wrap"}}>{commandInfo.description}</pre><br/>
-                <Divider />
-                <b>Requires Admin?</b><pre style={{margin:0}}>{commandInfo.needs_admin ? "True": "False"}</pre><br/>
-                <Divider />
-                {parameterGroups.length > 1 &&
-                    <FormControl style={{width: "100%", marginTop: "7px"}} >
+            <Box className="mythic-task-parameters-overview">
+                <Typography component="div" className="mythic-task-parameters-section-label">
+                    Description
+                </Typography>
+                <Typography component="pre" className="mythic-task-parameters-description">
+                    {commandInfo.description || "No description provided."}
+                </Typography>
+            </Box>
+            {parameterGroups.length > 1 &&
+                <Box className="mythic-task-parameters-group-card">
+                    <Box className="mythic-task-parameters-group-copy">
+                        <Typography component="div" className="mythic-task-parameters-section-label">
+                            Parameter group
+                        </Typography>
+                        <Typography component="div" className="mythic-task-parameters-section-description">
+                            {parameterGroups.length} available groups
+                        </Typography>
+                    </Box>
+                    <FormControl className="mythic-task-parameters-group-select">
                        <TextField
                            select
+                            size="small"
                            label="Parameter Group"
                            value={selectedParameterGroup}
                            onChange={onChangeParameterGroup}
-                            
-                            input={<Input />}
                        >
                        {
                            parameterGroups.map((opt, i) => (
@@ -1059,42 +1097,31 @@ export function TaskParametersDialog(props) {
                        }
                        </TextField>
                    </FormControl>
-                    
-                }
-            </Typography>
-            <TableContainer className="mythicElement">
-                <Table size="small" style={{"tableLayout": "fixed", "maxWidth": "100%", "overflow": "scroll"}}>
-                    <TableHead>
-                        <TableRow>
-                            <TableCell style={{width: "30%"}}>Parameter</TableCell>
-                            <TableCell>Value</TableCell>
-                        </TableRow>
-                    </TableHead>
-                    <TableBody>
-                        {parameters.map( (op) => (
-                            <TaskParametersDialogRow onSubmit={onSubmit} key={"taskparameterrow" + op.id}
-                                onChange={onChange} commandInfo={commandInfo} {...op}
-                                parameterGroupName={selectedParameterGroup}
-                                callback_id={props.callback_id}
-                                onAgentConnectAddNewPayloadOnHost={onAgentConnectAddNewPayloadOnHost}
-                                onAgentConnectRemovePayloadOnHost={onAgentConnectRemovePayloadOnHost}
-                                addedCredential={addedCredential} removedCredential={removedCredential}
-                                setSubmenuOpenPreventTasking={setSubmenuOpenPreventTasking}
-                                                     getOtherParameters={getOtherParameters}
-                                />
-                        ))}
-                    </TableBody>
-                </Table>
-            </TableContainer>
+                </Box>
+            }
+            <Box className="mythic-task-parameters-list">
+                {parameters.map( (op) => (
+                    <TaskParametersDialogRow onSubmit={onSubmit} key={"taskparameterrow" + op.id}
+                        onChange={onChange} commandInfo={commandInfo} {...op}
+                        parameterGroupName={selectedParameterGroup}
+                        callback_id={props.callback_id}
+                        onAgentConnectAddNewPayloadOnHost={onAgentConnectAddNewPayloadOnHost}
+                        onAgentConnectRemovePayloadOnHost={onAgentConnectRemovePayloadOnHost}
+                        addedCredential={addedCredential} removedCredential={removedCredential}
+                        setSubmenuOpenPreventTasking={setSubmenuOpenPreventTasking}
+                                             getOtherParameters={getOtherParameters}
+                    />
+                ))}
+            </Box>
        </DialogContent>
-        <DialogActions>
-          <Button onClick={props.onClose} variant="contained" color="primary">
+        <MythicDialogFooter className="mythic-task-parameters-actions">
+          <MythicDialogButton onClick={props.onClose}>
            Close
-          </Button>
-          <Button onClick={onSubmit} disabled={submenuOpenPreventTask} variant="contained" color="warning">
+          </MythicDialogButton>
+          <MythicDialogButton onClick={onSubmit} disabled={submenuOpenPreventTask} intent="primary">
            {props.captureOnly ? "Use Parameters" : "Task"}
-          </Button>
-        </DialogActions>
+          </MythicDialogButton>
+        </MythicDialogFooter>
  </React.Fragment>
  );
 }
@@ -341,32 +341,47 @@ tspan {
    width: 100%;
 }
 .mythic-tasking-context-row {
-    align-items: center;
+    align-items: flex-start;
    display: flex;
+    flex-wrap: wrap;
    gap: 0.28rem;
    margin-bottom: 0.45rem;
    min-height: 24px;
    min-width: 0;
-    overflow-x: auto;
-    overflow-y: hidden;
    padding-bottom: 0.05rem;
 }
+.mythic-tasking-context-row > span[data-tooltip-id="my-tooltip"] {
+    display: inline-flex !important;
+    max-width: 100%;
+    min-width: 0;
+}
 .mythic-tasking-context-chip {
-    align-items: center;
+    align-items: flex-start;
    border: 1px solid ${(props) => alpha(props.theme.palette.common.black, props.theme.palette.mode === "dark" ? 0.22 : 0.1)};
    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    cursor: pointer;
    display: inline-flex;
-    flex: 0 0 auto;
+    flex: 0 1 auto;
    font-size: 0.72rem;
    font-weight: 750;
    gap: 0.28rem;
-    line-height: 1;
-    max-width: min(28rem, 72vw);
+    line-height: 1.2;
+    max-width: 100%;
    min-height: 22px;
-    overflow: hidden;
-    padding: 0.22rem 0.45rem;
-    text-overflow: ellipsis;
-    white-space: nowrap;
+    overflow-wrap: anywhere;
+    padding: 0.24rem 0.48rem;
+    transition: border-color 120ms ease, box-shadow 120ms ease, transform 120ms ease;
+    white-space: normal;
+    word-break: break-word;
+}
+.mythic-tasking-context-chip:hover,
+.mythic-tasking-context-chip:focus-visible {
+    box-shadow: inset 0 0 0 1px currentColor;
+    outline: none;
+    transform: translateY(-1px);
+}
+.mythic-tasking-context-chip:active {
+    transform: translateY(0);
 }
 .mythic-tasking-context-chip-emphasized {
    box-shadow: inset 0 0 0 1px ${(props) => alpha(props.theme.palette.common.white, props.theme.palette.mode === "dark" ? 0.12 : 0.18)};
@@ -380,8 +395,8 @@ tspan {
 }
 .mythic-tasking-context-chip-value {
    min-width: 0;
-    overflow: hidden;
-    text-overflow: ellipsis;
+    overflow-wrap: anywhere;
+    word-break: break-word;
 }
 .mythic-tasking-command-row {
    align-items: flex-end;
@@ -707,6 +722,578 @@ tspan {
        width: 100% !important;
    }
 }
+.mythic-task-parameters-title.MuiDialogTitle-root {
+    background-image: ${getSectionHeaderGradient} !important;
+    border-bottom: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor} !important;
+    color: ${(props) => props.theme.pageHeaderText?.main || props.theme.palette.text.primary};
+    cursor: move;
+    padding: 1rem 1.15rem !important;
+    position: relative;
+}
+.mythic-task-parameters-title.MuiDialogTitle-root::before {
+    background-color: ${getSectionHeaderAccent};
+    bottom: 0;
+    box-shadow: 0 0 0 1px ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.2)};
+    content: "";
+    left: 0;
+    position: absolute;
+    top: 0;
+    width: 6px;
+}
+.mythic-task-parameters-title-row {
+    align-items: center;
+    display: flex;
+    gap: 0.8rem;
+    min-height: 42px;
+    min-width: 0;
+    padding-left: 0.55rem;
+}
+.mythic-task-parameters-title-icon {
+    align-items: center;
+    background-color: ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.12)};
+    border: 1px solid ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.24)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: inline-flex;
+    flex: 0 0 auto;
+    height: 32px;
+    justify-content: center;
+    width: 32px;
+}
+.mythic-task-parameters-title-copy {
+    flex: 1 1 auto;
+    min-width: 12rem;
+}
+.mythic-task-parameters-title-main {
+    color: ${(props) => props.theme.pageHeaderText?.main || props.theme.palette.text.primary};
+    font-size: 1rem;
+    font-weight: 850;
+    line-height: 1.15;
+    overflow-wrap: anywhere;
+}
+.mythic-task-parameters-title-subtitle {
+    color: ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.78)};
+    font-size: 0.76rem;
+    font-weight: 600;
+    line-height: 1.25;
+    margin-top: 0.15rem;
+}
+.mythic-task-parameters-title-meta {
+    align-items: center;
+    display: flex;
+    flex: 0 1 auto;
+    flex-wrap: wrap;
+    gap: 0.35rem;
+    justify-content: flex-end;
+    min-width: 0;
+}
+.mythic-task-parameters-title-chip.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.12)};
+    border: 1px solid ${(props) => alpha(props.theme.pageHeaderText?.main || props.theme.palette.text.primary, 0.24)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.pageHeaderText?.main || props.theme.palette.text.primary};
+    font-size: 0.72rem;
+    font-weight: 750;
+}
+.mythic-task-parameters-title-chip .MuiChip-icon {
+    color: inherit;
+    font-size: 1rem;
+}
+.mythic-task-parameters-title-chip-warning.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.22 : 0.14)};
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.52 : 0.36)};
+    color: ${(props) => props.theme.palette.warning.main};
+}
+.mythic-task-parameters-content.MuiDialogContent-root {
+    background-color: ${(props) => props.theme.surfaces?.raised || props.theme.palette.background.paper};
+    max-height: min(76vh, 52rem);
+    min-width: min(58rem, calc(100vw - 3rem));
+    overflow: auto;
+    padding: 1rem !important;
+    position: relative;
+}
+.mythic-task-parameters-backdrop.MuiBackdrop-root {
+    position: absolute;
+    z-index: 2;
+}
+.mythic-task-parameters-overview,
+.mythic-task-parameters-group-card,
+.mythic-task-parameter-card {
+    background-color: ${(props) => props.theme.surfaces?.muted || props.theme.palette.background.paper};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    min-width: 0;
+}
+.mythic-task-parameters-overview {
+    padding: 0.75rem;
+}
+.mythic-task-parameters-section-label {
+    color: ${(props) => props.theme.palette.text.primary};
+    font-size: 0.78rem;
+    font-weight: 800;
+    line-height: 1.25;
+}
+.mythic-task-parameters-section-description {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.74rem;
+    line-height: 1.35;
+    margin-top: 0.12rem;
+}
+.mythic-task-parameters-description {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-family: ${(props) => props.theme.typography.fontFamily};
+    font-size: 0.82rem;
+    line-height: 1.45;
+    margin: 0.35rem 0 0;
+    overflow-wrap: anywhere;
+    white-space: pre-wrap;
+}
+.mythic-task-parameters-group-card {
+    align-items: center;
+    display: grid;
+    gap: 0.75rem;
+    grid-template-columns: minmax(0, 1fr) minmax(14rem, 20rem);
+    padding: 0.75rem;
+}
+.mythic-task-parameters-group-select {
+    width: 100%;
+}
+.mythic-task-parameters-list {
+    display: flex;
+    flex-direction: column;
+    gap: 0.7rem;
+    margin-top: 0.75rem;
+    min-width: 0;
+}
+.mythic-task-parameter-card {
+    display: grid;
+    gap: 0.85rem;
+    grid-template-columns: minmax(14rem, 0.38fr) minmax(0, 0.62fr);
+    padding: 0.8rem;
+}
+.mythic-task-parameter-card-required {
+    border-left: 4px solid ${(props) => props.theme.palette.warning.main};
+}
+.mythic-task-parameter-copy {
+    min-width: 0;
+}
+.mythic-task-parameter-heading {
+    align-items: flex-start;
+    display: flex;
+    gap: 0.5rem;
+    justify-content: space-between;
+    min-width: 0;
+}
+.mythic-task-parameter-title {
+    color: ${(props) => props.theme.palette.text.primary};
+    font-size: 0.88rem;
+    font-weight: 800;
+    line-height: 1.25;
+    min-width: 0;
+    overflow-wrap: anywhere;
+}
+.mythic-task-parameter-chip-row {
+    align-items: center;
+    display: flex;
+    flex: 0 0 auto;
+    flex-wrap: wrap;
+    gap: 0.25rem;
+    justify-content: flex-end;
+}
+.mythic-task-parameter-chip.MuiChip-root {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.055)" : "rgba(0,0,0,0.032)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.68rem;
+    font-weight: 750;
+    height: 22px;
+}
+.mythic-task-parameter-chip-required.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.2 : 0.12)};
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.5 : 0.34)};
+    color: ${(props) => props.theme.palette.warning.main};
+}
+.mythic-task-parameter-description {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.76rem;
+    line-height: 1.4;
+    margin-top: 0.35rem;
+    overflow-wrap: anywhere;
+}
+.mythic-task-parameter-description-muted {
+    color: ${(props) => props.theme.palette.text.disabled};
+    font-style: italic;
+}
+.mythic-task-parameter-name {
+    color: ${(props) => props.theme.palette.text.disabled};
+    font-family: ${(props) => props.theme.typography.fontFamilyMono || "monospace"};
+    font-size: 0.7rem;
+    line-height: 1.35;
+    margin-top: 0.4rem;
+    overflow-wrap: anywhere;
+}
+.mythic-task-parameter-control {
+    min-width: 0;
+    position: relative;
+}
+.mythic-task-parameter-control-shell {
+    min-width: 0;
+    position: relative;
+}
+.mythic-task-parameter-control-backdrop.MuiBackdrop-root {
+    background-color: ${(props) => alpha(props.theme.palette.background.paper, props.theme.palette.mode === "dark" ? 0.72 : 0.82)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.primary};
+    position: absolute;
+    z-index: 2;
+}
+.mythic-task-parameter-control .MuiFormControl-root,
+.mythic-task-parameter-control .MuiTextField-root {
+    width: 100%;
+}
+.mythic-task-parameter-control .MuiInputBase-root {
+    min-height: 38px;
+}
+.mythic-task-parameter-control .MuiInput-root {
+    background-color: ${(props) => props.theme.palette.background.paper};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    padding-left: 0.45rem;
+    padding-right: 0.45rem;
+}
+.mythic-task-parameter-control .MuiInput-root::before,
+.mythic-task-parameter-control .MuiInput-root::after {
+    display: none;
+}
+.mythic-task-parameter-control .mythicElement {
+    background-color: transparent;
+    border: 0;
+    box-shadow: none;
+}
+.mythic-task-parameter-select-row {
+    align-items: center;
+    display: grid;
+    gap: 0.45rem;
+    grid-template-columns: minmax(0, 1fr) auto;
+    min-width: 0;
+}
+.mythic-task-parameter-select-control {
+    min-width: 0;
+    width: 100%;
+}
+.mythic-task-parameter-select-value,
+.mythic-task-parameter-select-placeholder {
+    display: block;
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+.mythic-task-parameter-select-placeholder {
+    color: ${(props) => props.theme.palette.text.disabled};
+}
+.mythic-task-parameter-selected-values {
+    align-items: center;
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.25rem;
+    min-width: 0;
+}
+.mythic-task-parameter-selected-chip.MuiChip-root {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.06)" : "rgba(0,0,0,0.035)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.68rem;
+    height: 22px;
+    max-width: 100%;
+}
+.mythic-task-parameter-selected-chip .MuiChip-label {
+    overflow: hidden;
+    text-overflow: ellipsis;
+}
+.mythic-task-parameter-menu-text {
+    display: inline-block;
+    min-width: 0;
+    overflow-wrap: anywhere;
+    white-space: pre-wrap;
+    word-break: break-word;
+}
+.mythic-task-parameter-empty-inline {
+    align-items: center;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.035)" : "rgba(0,0,0,0.018)"};
+    border: 1px dashed ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.secondary};
+    display: flex;
+    font-size: 0.74rem;
+    font-weight: 650;
+    min-height: 38px;
+    padding: 0.45rem 0.65rem;
+}
+.mythic-task-parameter-refresh.MuiIconButton-root {
+    align-self: center;
+}
+.mythic-task-choice-custom {
+    align-items: center;
+    display: grid;
+    gap: 0.55rem;
+    grid-template-columns: minmax(0, 1fr) auto minmax(0, 1fr) auto;
+    min-width: 0;
+}
+.mythic-task-choice-custom-select,
+.mythic-task-choice-custom-input {
+    min-width: 0;
+}
+.mythic-task-choice-custom-divider {
+    align-items: center;
+    color: ${(props) => props.theme.palette.text.disabled};
+    display: inline-flex;
+    font-size: 0.68rem;
+    font-weight: 800;
+    justify-content: center;
+}
+.mythic-task-choice-custom-divider span {
+    background-color: ${(props) => props.theme.surfaces?.muted || props.theme.palette.background.paper};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: 999px;
+    padding: 0.16rem 0.4rem;
+}
+.mythic-task-parameter-boolean-row {
+    align-items: center;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.035)" : "rgba(0,0,0,0.018)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: inline-flex;
+    gap: 0.45rem;
+    padding: 0.25rem 0.55rem 0.25rem 0.35rem;
+}
+.mythic-task-parameter-boolean-chip.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border: 1px solid ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.warning.main};
+    font-size: 0.68rem;
+    font-weight: 750;
+    height: 22px;
+}
+.mythic-task-parameter-boolean-chip-enabled.MuiChip-root {
+    background-color: ${(props) => alpha(props.theme.palette.success.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border-color: ${(props) => alpha(props.theme.palette.success.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)};
+    color: ${(props) => props.theme.palette.success.main};
+}
+.mythic-task-parameter-field-row {
+    align-items: center;
+    display: grid;
+    gap: 0.65rem;
+    grid-template-columns: minmax(8.5rem, 0.28fr) minmax(0, 0.72fr);
+    min-width: 0;
+}
+.mythic-task-parameter-field-label {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.72rem;
+    font-weight: 750;
+    line-height: 1.3;
+    min-width: 0;
+    overflow-wrap: anywhere;
+}
+.mythic-task-parameter-field-control {
+    min-width: 0;
+}
+.mythic-task-parameter-field-actions {
+    align-items: center;
+    display: flex;
+    gap: 0.4rem;
+    justify-content: flex-end;
+}
+.mythic-task-parameter-switch-row {
+    align-items: center;
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.035)" : "rgba(0,0,0,0.018)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    display: flex;
+    gap: 0.6rem;
+    justify-content: space-between;
+    padding: 0.45rem 0.65rem;
+}
+.mythic-task-array-editor,
+.mythic-task-agent-connect-editor {
+    display: flex;
+    flex-direction: column;
+    gap: 0.65rem;
+    min-width: 0;
+}
+.mythic-task-array-list,
+.mythic-task-agent-connect-panel,
+.mythic-task-agent-connect-parameter-list {
+    display: flex;
+    flex-direction: column;
+    gap: 0.55rem;
+    min-width: 0;
+}
+.mythic-task-array-entry,
+.mythic-task-typed-array-entry {
+    align-items: center;
+    display: grid;
+    gap: 0.55rem;
+    min-width: 0;
+}
+.mythic-task-array-entry {
+    grid-template-columns: auto minmax(0, 1fr);
+}
+.mythic-task-typed-array-entry {
+    grid-template-columns: auto minmax(8rem, 0.3fr) minmax(0, 0.7fr);
+}
+.mythic-task-array-delete.MuiIconButton-root {
+    align-self: center;
+}
+.mythic-task-array-entry-control,
+.mythic-task-typed-array-value,
+.mythic-task-typed-array-choice {
+    min-width: 0;
+}
+.mythic-task-parameter-add-button.MuiButton-root {
+    align-self: flex-start;
+}
+.mythic-task-agent-connect-panel,
+.mythic-task-agent-connect-parameters {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.035)" : "rgba(0,0,0,0.018)"};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    padding: 0.65rem;
+}
+.mythic-task-agent-connect-actions {
+    align-items: center;
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.45rem;
+    justify-content: flex-end;
+}
+.mythic-task-agent-connect-parameters-title {
+    color: ${(props) => props.theme.palette.text.primary};
+    font-size: 0.78rem;
+    font-weight: 800;
+    line-height: 1.3;
+    margin-bottom: 0.55rem;
+}
+.mythic-task-agent-connect-parameter {
+    align-items: start;
+    display: grid;
+    gap: 0.65rem;
+    grid-template-columns: minmax(8rem, 0.32fr) minmax(0, 0.68fr);
+    min-width: 0;
+}
+.mythic-task-agent-connect-parameter-name {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-family: ${(props) => props.theme.typography.fontFamilyMono || "monospace"};
+    font-size: 0.72rem;
+    line-height: 1.35;
+    overflow-wrap: anywhere;
+}
+.mythic-task-agent-connect-parameter-value {
+    background-color: ${(props) => props.theme.palette.background.paper};
+    border: 1px solid ${(props) => props.theme.table?.borderSoft || props.theme.borderColor};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    color: ${(props) => props.theme.palette.text.primary};
+    font-family: ${(props) => props.theme.typography.fontFamilyMono || "monospace"};
+    font-size: 0.72rem;
+    line-height: 1.4;
+    margin: 0;
+    max-height: 14rem;
+    min-width: 0;
+    overflow: auto;
+    padding: 0.5rem;
+    white-space: pre-wrap;
+    word-break: break-word;
+}
+.mythic-task-credential-editor {
+    display: flex;
+    flex-direction: column;
+    gap: 0.55rem;
+    min-width: 0;
+}
+.mythic-task-credential-actions {
+    align-items: center;
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.45rem;
+    justify-content: flex-end;
+}
+.mythic-task-credential-menu-item {
+    display: flex;
+    flex-direction: column;
+    gap: 0.18rem;
+    min-width: 0;
+}
+.mythic-task-credential-menu-comment {
+    color: ${(props) => props.theme.palette.text.secondary};
+    font-size: 0.72rem;
+    line-height: 1.3;
+    overflow-wrap: anywhere;
+    white-space: pre-wrap;
+}
+.mythic-task-file-dropzone {
+    background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.035)" : "rgba(0,0,0,0.018)"};
+    border: 1px dashed ${(props) => alpha(props.theme.palette.primary.main, props.theme.palette.mode === "dark" ? 0.58 : 0.42)};
+    border-radius: ${(props) => props.theme.shape.borderRadius}px;
+    cursor: pointer;
+    padding: 1rem;
+    text-align: center;
+    transition: background-color 120ms ease, border-color 120ms ease, box-shadow 120ms ease;
+}
+.mythic-task-file-dropzone:hover,
+.mythic-task-file-dropzone-dragging {
+    background-color: ${(props) => alpha(props.theme.palette.primary.main, props.theme.palette.mode === "dark" ? 0.12 : 0.06)};
+    border-color: ${(props) => props.theme.palette.primary.main};
+    box-shadow: inset 0 0 0 1px ${(props) => alpha(props.theme.palette.primary.main, 0.28)};
+}
+.mythic-task-file-dropzone-content {
+    align-items: center;
+    color: ${(props) => props.theme.palette.text.secondary};
+    display: flex;
+    flex-direction: column;
+    gap: 0.2rem;
+    min-width: 0;
+    overflow-wrap: anywhere;
+}
+.mythic-task-file-dropzone-icon {
+    color: ${(props) => props.theme.palette.primary.main};
+}
+.mythic-task-parameters-actions.MuiDialogActions-root {
+    padding: 0.85rem 1rem !important;
+}
+@media screen and (max-width: 900px) {
+    .mythic-task-parameters-title-row,
+    .mythic-task-parameters-group-card,
+    .mythic-task-parameter-card {
+        grid-template-columns: 1fr;
+    }
+    .mythic-task-parameters-title-row {
+        align-items: flex-start;
+        flex-wrap: wrap;
+    }
+    .mythic-task-parameters-title-meta {
+        justify-content: flex-start;
+        width: 100%;
+    }
+    .mythic-task-parameters-content.MuiDialogContent-root {
+        min-width: 0;
+    }
+    .mythic-task-parameter-field-row,
+    .mythic-task-typed-array-entry,
+    .mythic-task-agent-connect-parameter {
+        grid-template-columns: 1fr;
+    }
+    .mythic-task-choice-custom {
+        grid-template-columns: 1fr;
+    }
+    .mythic-task-choice-custom-divider {
+        justify-content: flex-start;
+    }
+    .mythic-task-array-entry {
+        align-items: flex-start;
+    }
+}
 .mythic-file-browser-iconButtonCompound.MuiIconButton-root {
    gap: 0;
    padding: 0 0.15rem;
@@ -47,6 +47,8 @@ type AgentMessageRawInput struct {
 }

 type cachedUUIDInfo struct {
+	stateMutex                   sync.RWMutex
+	cryptoMutex                  sync.Mutex
 	UUID                         string
 	UUIDType                     string // payload, callback, staging
 	C2ProfileName                string
@@ -140,31 +142,39 @@ func (cache *cachedUUIDInfo) IterateAndAct(agentMessage *[]byte, action string)
 		case "aes256_hmac":
 			if action == "decrypt" {
 				keyToUse := keyOpt.DecKey
+				cache.cryptoMutex.Lock()
 				if cache.SuccessfulDecKeyOpt != nil {
 					keyToUse = cache.SuccessfulDecKeyOpt
 				}
+				cache.cryptoMutex.Unlock()
 				if modified, err = mythicCrypto.DecryptAES256HMAC(*keyToUse, agentMessage); err == nil {
 					// we successfully decrypted, so return
 					// track the key actually used so that when we encrypt for the return message we can use the right one
 					//logging.LogDebug("setting uuidInfo successfulDecKeyOpt")
+					cache.cryptoMutex.Lock()
 					cache.SuccessfulDecKeyOpt = keyToUse
 					cache.SuccessfulEncKeyOpt = keyToUse
+					cache.cryptoMutex.Unlock()
 					//logging.LogDebug("just updated uuindInfo", "updated cache", cache)
 					return modified, err
 				}
 			} else if action == "encrypt" {
 				//logging.LogDebug("encrypting", "current cache", cache)
 				keyToUse := keyOpt.EncKey
+				cache.cryptoMutex.Lock()
 				if cache.SuccessfulEncKeyOpt != nil {
 					keyToUse = cache.SuccessfulEncKeyOpt
 				}
+				cache.cryptoMutex.Unlock()
 				//logging.LogDebug("Encrypting message", "key", hex.EncodeToString(*keyToUse))

 				modified, err = mythicCrypto.EncryptAES256HMAC(*keyToUse, agentMessage)
 				if err == nil {
 					// we successfully encrypted, so return
+					cache.cryptoMutex.Lock()
 					cache.SuccessfulDecKeyOpt = keyToUse
 					cache.SuccessfulEncKeyOpt = keyToUse
+					cache.cryptoMutex.Unlock()
 					//logging.LogDebug("Encrypted message", "output", hex.EncodeToString(modified))
 					return modified, err
 				}
@@ -186,7 +196,7 @@ func (cache *cachedUUIDInfo) IterateAndAct(agentMessage *[]byte, action string)
 	return modified, err
 }

-var cachedUUIDInfoMapMutex = &sync.Mutex{}
+var cachedUUIDInfoMapMutex = &sync.RWMutex{}

 func InvalidateAllCachedUUIDInfo() {
 	cachedUUIDInfoMapMutex.Lock()
@@ -206,22 +216,63 @@ func InvalidateCachedUUIDInfo(uuid string) {

 var cachedUUIDInfoMap = make(map[string]*cachedUUIDInfo)

+func getCachedUUIDInfo(c2profile string, messageUUID string) *cachedUUIDInfo {
+	cachedUUIDInfoMapMutex.RLock()
+	defer cachedUUIDInfoMapMutex.RUnlock()
+	if uuidInfo, ok := cachedUUIDInfoMap[messageUUID+c2profile]; ok {
+		return uuidInfo
+	}
+	if uuidInfo, ok := cachedUUIDInfoMap[messageUUID]; ok {
+		return uuidInfo
+	}
+	return nil
+}
+
+func cacheUUIDInfo(c2profile string, messageUUID string, newCache *cachedUUIDInfo) (*cachedUUIDInfo, bool) {
+	cachedUUIDInfoMapMutex.Lock()
+	defer cachedUUIDInfoMapMutex.Unlock()
+	if uuidInfo, ok := cachedUUIDInfoMap[messageUUID+c2profile]; ok {
+		return uuidInfo, true
+	}
+	if uuidInfo, ok := cachedUUIDInfoMap[messageUUID]; ok {
+		return uuidInfo, true
+	}
+	cachedUUIDInfoMap[messageUUID+c2profile] = newCache
+	return newCache, false
+}
+
+func updateCachedCallbackCheckin(uuidInfo *cachedUUIDInfo, updateCheckinTime bool) {
+	if updateCheckinTime && uuidInfo.UUIDType == UUIDTYPECALLBACK {
+		UpdateCallbackEdgesAndCheckinTime(uuidInfo)
+	}
+}
+
+func (cache *cachedUUIDInfo) callbackAllowedFromPayload() bool {
+	cache.stateMutex.RLock()
+	defer cache.stateMutex.RUnlock()
+	return cache.CallbackAllowedFromPayload
+}
+
 func MarkCallbackInfoInactive(callbackID int) {
 	cachedUUIDInfoMapMutex.Lock()
 	defer cachedUUIDInfoMapMutex.Unlock()
-	for k, _ := range cachedUUIDInfoMap {
-		if cachedUUIDInfoMap[k].CallbackID == callbackID {
-			cachedUUIDInfoMap[k].Active = false
-			cachedUUIDInfoMap[k].EdgeId = 0
+	for _, uuidInfo := range cachedUUIDInfoMap {
+		if uuidInfo.CallbackID == callbackID {
+			uuidInfo.stateMutex.Lock()
+			uuidInfo.Active = false
+			uuidInfo.EdgeId = 0
+			uuidInfo.stateMutex.Unlock()
 		}
 	}
 }
 func UpdateCallbackInfoTriggerOnCheckin(callbackID int, triggerOnCheckinAfterTime int) {
 	cachedUUIDInfoMapMutex.Lock()
 	defer cachedUUIDInfoMapMutex.Unlock()
-	for k, _ := range cachedUUIDInfoMap {
-		if cachedUUIDInfoMap[k].CallbackID == callbackID {
-			cachedUUIDInfoMap[k].TriggerOnCheckinAfterTime = triggerOnCheckinAfterTime
+	for _, uuidInfo := range cachedUUIDInfoMap {
+		if uuidInfo.CallbackID == callbackID {
+			uuidInfo.stateMutex.Lock()
+			uuidInfo.TriggerOnCheckinAfterTime = triggerOnCheckinAfterTime
+			uuidInfo.stateMutex.Unlock()
 			return
 		}
 	}
@@ -229,9 +280,11 @@ func UpdateCallbackInfoTriggerOnCheckin(callbackID int, triggerOnCheckinAfterTim
 func UpdatePayloadInfoCallbackAllowed(payloadID int, callbackAllowed bool) {
 	cachedUUIDInfoMapMutex.Lock()
 	defer cachedUUIDInfoMapMutex.Unlock()
-	for k, _ := range cachedUUIDInfoMap {
-		if cachedUUIDInfoMap[k].PayloadID == payloadID && cachedUUIDInfoMap[k].UUIDType == UUIDTYPEPAYLOAD {
-			cachedUUIDInfoMap[k].CallbackAllowedFromPayload = callbackAllowed
+	for _, uuidInfo := range cachedUUIDInfoMap {
+		if uuidInfo.PayloadID == payloadID && uuidInfo.UUIDType == UUIDTYPEPAYLOAD {
+			uuidInfo.stateMutex.Lock()
+			uuidInfo.CallbackAllowedFromPayload = callbackAllowed
+			uuidInfo.stateMutex.Unlock()
 			return
 		}
 	}
@@ -349,23 +402,6 @@ func processAgentMessageContent(agentMessageInput *AgentMessageRawInput, uuidInf
 				go SendAllOperationsMessage(err.Error(), uuidInfo.OperationID, "agent_message_bad_get_tasking", database.MESSAGE_LEVEL_AGENT_MESSGAGE, true)
 			}
 		}
-	case "upload":
-		{
-			go SendAllOperationsMessage(fmt.Sprintf("Agent %s is using deprecated method of file transfer with the 'upload' action.", uuidInfo.PayloadTypeName),
-				uuidInfo.OperationID, "debug", database.MESSAGE_LEVEL_AGENT_MESSGAGE, false)
-			logging.LogError(nil, "deprecated form of upload detected, please update agent code to use new method")
-			response, err = handleAgentMessagePostResponse(&map[string]interface{}{
-				"action": "post_response",
-				"responses": []map[string]interface{}{
-					{
-						"task_id": decryptedMessage["task_id"],
-						"upload":  decryptedMessage,
-					},
-				},
-			}, uuidInfo)
-			uploadResponse := response["responses"].([]map[string]interface{})
-			response = uploadResponse[0]
-		}
 	case "post_response":
 		{
 			response, err = handleAgentMessagePostResponse(&decryptedMessage, uuidInfo)
@@ -533,20 +569,10 @@ func processAgentMessageContent(agentMessageInput *AgentMessageRawInput, uuidInf
 		response["delegates"] = delegateResponses
 	}
 	// get first order proxy data not for delegate callbacks
-	if proxyData, err := proxyPorts.GetDataForCallbackIdPortType(uuidInfo.CallbackID, CALLBACK_PORT_TYPE_SOCKS); err != nil {
+	if proxyData, err := proxyPorts.GetDataForCallbackIdAllTypes(uuidInfo.CallbackID); err != nil {
 		logging.LogError(err, "Failed to get proxy data")
-	} else if proxyData != nil {
-		response[CALLBACK_PORT_TYPE_SOCKS] = proxyData
-	}
-	if proxyData, err := proxyPorts.GetDataForCallbackIdPortType(uuidInfo.CallbackID, CALLBACK_PORT_TYPE_RPORTFWD); err != nil {
-		logging.LogError(err, "Failed to get proxy data")
-	} else if proxyData != nil {
-		response[CALLBACK_PORT_TYPE_RPORTFWD] = proxyData
-	}
-	if proxyData, err := proxyPorts.GetDataForCallbackIdPortType(uuidInfo.CallbackID, CALLBACK_PORT_TYPE_INTERACTIVE); err != nil {
-		logging.LogError(err, "Failed to get interactive data")
-	} else if proxyData != nil {
-		response[CALLBACK_PORT_TYPE_INTERACTIVE] = proxyData
+	} else {
+		proxyData.AddToResponse(response)
 	}
 	response["action"] = decryptedMessage["action"]
 	// reflect back any non-standard key at the top level
@@ -679,7 +705,7 @@ func recursiveProcessAgentMessage(agentMessageInput *AgentMessageRawInput) recur
 		instanceResponse.Err = errors.New(errorMessage)
 		return instanceResponse
 	}
-	if !uuidInfo.CallbackAllowedFromPayload {
+	if !uuidInfo.callbackAllowedFromPayload() {
 		payload := databaseStructs.Payload{}
 		err = database.DB.Get(&payload, `SELECT
 			payload.description, 
@@ -735,26 +761,9 @@ func LookupEncryptionData(c2profile string, messageUUID string, updateCheckinTim
 	//logging.LogDebug("Getting encryption data", "cachemap", cachedUUIDInfoMap)
 	newCache := cachedUUIDInfo{}

-	cachedUUIDInfoMapMutex.Lock()
-	defer cachedUUIDInfoMapMutex.Unlock()
-	if _, ok := cachedUUIDInfoMap[messageUUID+c2profile]; ok {
-		// we found an instance of the cache info with c2 profile encryption data
-		if cachedUUIDInfoMap[messageUUID+c2profile].UUIDType == UUIDTYPECALLBACK {
-			if updateCheckinTime {
-				UpdateCallbackEdgesAndCheckinTime(cachedUUIDInfoMap[messageUUID+c2profile])
-			}
-
-		}
-		return cachedUUIDInfoMap[messageUUID+c2profile], nil
-	} else if _, ok = cachedUUIDInfoMap[messageUUID]; ok {
-		// we found an instance of the cache info with payload encryption data
-		if cachedUUIDInfoMap[messageUUID].UUIDType == UUIDTYPECALLBACK {
-			if updateCheckinTime {
-				UpdateCallbackEdgesAndCheckinTime(cachedUUIDInfoMap[messageUUID])
-			}
-
-		}
-		return cachedUUIDInfoMap[messageUUID], nil
+	if uuidInfo := getCachedUUIDInfo(c2profile, messageUUID); uuidInfo != nil {
+		updateCachedCallbackCheckin(uuidInfo, updateCheckinTime)
+		return uuidInfo, nil
 	}

 	// get the associated c2 profile
@@ -944,9 +953,13 @@ func LookupEncryptionData(c2profile string, messageUUID string, updateCheckinTim
 			return &newCache, nil
 		}
 	}
-	cachedUUIDInfoMap[messageUUID+c2profile] = &newCache
-	logging.LogDebug("New cache value for agent connection", "newcache", newCache)
-	return &newCache, nil
+	uuidInfo, alreadyCached := cacheUUIDInfo(c2profile, messageUUID, &newCache)
+	if alreadyCached {
+		updateCachedCallbackCheckin(uuidInfo, updateCheckinTime)
+	} else {
+		logging.LogDebug("New cache value for agent connection", "newcache", uuidInfo)
+	}
+	return uuidInfo, nil
 }

 func DecryptMessage(uuidInfo *cachedUUIDInfo, agentMessage *[]byte) (map[string]interface{}, error) {
@@ -1272,6 +1285,9 @@ func updateCheckinTimeEverySecond() {
 	}
 }
 func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
+	uuidInfo.stateMutex.Lock()
+	defer uuidInfo.stateMutex.Unlock()
+
 	callback := databaseStructs.Callback{
 		AgentCallbackID: uuidInfo.UUID,
 		ID:              uuidInfo.CallbackID,
@@ -1280,7 +1296,12 @@ func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
 	}
 	// only bother updating the last checkin time if it's been more than one second
 	if callback.LastCheckin.Sub(uuidInfo.LastCheckinTime).Seconds() > 1 {
-		updateCheckinTimeChannel <- uuidInfo
+		previousCheckin := uuidInfo.LastCheckinTime
+		select {
+		case updateCheckinTimeChannel <- uuidInfo:
+		default:
+			logging.LogDebug("Skipping callback checkin update enqueue because channel is full", "callback_id", uuidInfo.CallbackID)
+		}
 		//callbackGraph.AddByAgentIds(callback.AgentCallbackID, callback.AgentCallbackID, uuidInfo.C2ProfileName)
 		if uuidInfo.EdgeId == 0 {
 			err := database.DB.Get(&uuidInfo.EdgeId, `SELECT id FROM callbackgraphedge
@@ -1288,9 +1309,10 @@ func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
 				uuidInfo.CallbackID, uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
 			if errors.Is(err, sql.ErrNoRows) {
 				if !uuidInfo.IsP2P {
-					_, err = database.DB.Exec(`INSERT INTO callbackgraphedge
+					err = database.DB.Get(&uuidInfo.EdgeId, `INSERT INTO callbackgraphedge
 						(source_id, destination_id, c2_profile_id, operation_id)
-						VALUES ($1, $1, $2, $3)`,
+						VALUES ($1, $1, $2, $3)
+						RETURNING id`,
 						uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
 					if err != nil {
 						logging.LogError(err, "Failed to add callback graph edge id for callback checking in",
@@ -1322,7 +1344,7 @@ func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
 			}
 		}
 		if uuidInfo.TriggerOnCheckinAfterTime > 0 {
-			checkinDifference := int(callback.LastCheckin.Sub(uuidInfo.LastCheckinTime).Minutes())
+			checkinDifference := int(callback.LastCheckin.Sub(previousCheckin).Minutes())
 			if checkinDifference >= uuidInfo.TriggerOnCheckinAfterTime {
 				// we want to trigger a workflow that the callback is checking in again after sleeping for > some time
 				go func(triggerData databaseStructs.Callback, oldCheckin time.Time, difference int) {
@@ -1335,7 +1357,7 @@ func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
 							"checkin_difference": difference,
 						},
 					}
-				}(callback, uuidInfo.LastCheckinTime, checkinDifference)
+				}(callback, previousCheckin, checkinDifference)
 			}
 		}
 		uuidInfo.LastCheckinTime = callback.LastCheckin
@@ -1,14 +1,15 @@
 package rabbitmq

 import (
+	"database/sql"
 	"errors"
 	"fmt"
-	"github.com/jmoiron/sqlx"
 	"time"

 	"github.com/its-a-feature/Mythic/database"
 	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
 	"github.com/its-a-feature/Mythic/logging"
+	"github.com/jmoiron/sqlx"
 	"github.com/mitchellh/mapstructure"
 )

@@ -26,6 +27,80 @@ type agentMessageGetTaskingTask struct {
 	Token      *int   `json:"token,omitempty"`
 }

+type agentMessageTaskRow struct {
+	ID           int           `db:"id"`
+	CallbackID   int           `db:"callback_id"`
+	AgentTaskID  string        `db:"agent_task_id"`
+	Timestamp    time.Time     `db:"timestamp"`
+	CommandName  string        `db:"command_name"`
+	Params       string        `db:"params"`
+	AgentTokenID sql.NullInt64 `db:"agent_token_id"`
+}
+
+func getAgentMessageTaskRows(taskIDs []int) ([]agentMessageTaskRow, error) {
+	currentTasks := []agentMessageTaskRow{}
+	if len(taskIDs) == 0 {
+		return currentTasks, nil
+	}
+	query, args, err := sqlx.Named(`SELECT
+		task.agent_task_id, task."timestamp", task.command_name, task.params, task.id, task.callback_id,
+		token.token_id "agent_token_id"
+		FROM task
+		LEFT JOIN token ON task.token_id = token.id
+		WHERE task.id IN (:ids)
+		ORDER BY task.id ASC`, map[string]interface{}{
+		"ids": taskIDs,
+	})
+	if err != nil {
+		return currentTasks, err
+	}
+	query, args, err = sqlx.In(query, args...)
+	if err != nil {
+		return currentTasks, err
+	}
+	query = database.DB.Rebind(query)
+	err = database.DB.Select(&currentTasks, query, args...)
+	return currentTasks, err
+}
+
+func markAgentMessageTasksProcessing(taskIDs []int, processingTime time.Time) error {
+	if len(taskIDs) == 0 {
+		return nil
+	}
+	query, args, err := sqlx.Named(`UPDATE task SET
+		status=:status,
+		status_timestamp_processing=:status_timestamp_processing
+		WHERE id IN (:ids)`, map[string]interface{}{
+		"ids":                         taskIDs,
+		"status":                      PT_TASK_FUNCTION_STATUS_PROCESSING,
+		"status_timestamp_processing": processingTime,
+	})
+	if err != nil {
+		return err
+	}
+	query, args, err = sqlx.In(query, args...)
+	if err != nil {
+		return err
+	}
+	query = database.DB.Rebind(query)
+	_, err = database.DB.Exec(query, args...)
+	return err
+}
+
+func buildAgentMessageTask(task agentMessageTaskRow) agentMessageGetTaskingTask {
+	newTask := agentMessageGetTaskingTask{
+		Command:    task.CommandName,
+		Parameters: task.Params,
+		ID:         task.AgentTaskID,
+		Timestamp:  task.Timestamp.Unix(),
+	}
+	if task.AgentTokenID.Valid {
+		tokenID := int(task.AgentTokenID.Int64)
+		newTask.Token = &tokenID
+	}
+	return newTask
+}
+
 func handleAgentMessageGetTasking(incoming *map[string]interface{}, callbackID int) (map[string]interface{}, error) {
 	// got message:
 	/*
@@ -118,80 +193,73 @@ func getDelegateTaskMessages(callbackID int, agentUUIDLength int, updateCheckinT
 	// get a list of all the other callbacks with tasks waiting to be processed
 	if callbackIds := submittedTasksAwaitingFetching.getOtherCallbackIds(callbackID); len(callbackIds) > 0 {
 		// check if there's a route between our callback and the callback with a task
+		routableCallbackIds := make([]int, 0, len(callbackIds))
+		routablePaths := make(map[int][]cbGraphAdjMatrixEntry)
 		for _, targetCallbackId := range callbackIds {
 			if routablePath := callbackGraph.GetBFSPath(callbackID, targetCallbackId); routablePath != nil && len(routablePath) > 0 {
 				// there's a route between our callback and the target callback for some sort of task
 				logging.LogDebug("task exists for callback we can route to")
-				currentTasks := []databaseStructs.Task{}
-				if taskIDs := submittedTasksAwaitingFetching.getTasksForCallbackId(targetCallbackId); len(taskIDs) > 0 {
-					query, args, err := sqlx.Named(`SELECT 
-						agent_task_id, "timestamp", command_name, params, id, token_id
-						FROM task WHERE id IN (:ids) ORDER BY id ASC`, map[string]interface{}{
-						"ids": taskIDs,
-					})
-					if err != nil {
-						logging.LogError(err, "Failed to make named statement when searching for tasks")
-						continue
-					}
-					query, args, err = sqlx.In(query, args...)
-					if err != nil {
-						logging.LogError(err, "Failed to do sqlx.In")
-						continue
-					}
-					query = database.DB.Rebind(query)
-					err = database.DB.Select(&currentTasks, query, args...)
-					if err != nil {
-						logging.LogError(err, "Failed to exec sqlx.IN modified statement")
-						continue
-					}
-					tasks := []agentMessageGetTaskingTask{}
-					for i := 0; i < len(currentTasks); i++ {
-						// now that we have a path, need to recursively encrypt and wrap
-						tokenID := 0
-						if currentTasks[i].TokenID.Valid {
-							if err := database.DB.Get(&tokenID, `SELECT token_id FROM token WHERE id=$1`, currentTasks[i].TokenID.Int64); err != nil {
-								logging.LogError(err, "failed to get token information")
-							}
-						}
-						newTask := agentMessageGetTaskingTask{
-							Command:    currentTasks[i].CommandName,
-							Parameters: currentTasks[i].Params,
-							ID:         currentTasks[i].AgentTaskID,
-							Timestamp:  currentTasks[i].Timestamp.Unix(),
-						}
-						if tokenID != 0 {
-							newTask.Token = &tokenID
-						}
-						_, err := database.DB.Exec(`UPDATE task SET
-							status=$2, status_timestamp_processing=$3
-							WHERE id=$1`, currentTasks[i].ID, PT_TASK_FUNCTION_STATUS_PROCESSING, time.Now().UTC())
-						if err != nil {
-							logging.LogError(err, "Failed to update task status to processing")
-							continue
-						}
-						tasks = append(tasks, newTask)
-						submittedTasksAwaitingFetching.removeTask(currentTasks[i].ID)
-						go addMitreAttackTaskMapping(currentTasks[i].ID)
-					}
-					if len(tasks) > 0 {
-						newTask := map[string]interface{}{
-							"action": "get_tasking",
-							"tasks":  tasks,
-						}
-						wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime)
-						if err != nil {
-							logging.LogError(err, "Failed to recursively encrypt message")
-							continue
-						}
-						delegateMessages = append(delegateMessages, delegateMessageResponse{
-							Message:       string(wrappedMessage),
-							SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
-							C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
-						})
-					}
-				}
+				routableCallbackIds = append(routableCallbackIds, targetCallbackId)
+				routablePaths[targetCallbackId] = routablePath
 			}
 		}
+		if len(routableCallbackIds) == 0 {
+			return delegateMessages
+		}
+		taskIDsByCallbackID := submittedTasksAwaitingFetching.getTasksForCallbackIds(routableCallbackIds)
+		allTaskIDs := make([]int, 0)
+		for _, targetCallbackId := range routableCallbackIds {
+			allTaskIDs = append(allTaskIDs, taskIDsByCallbackID[targetCallbackId]...)
+		}
+		if len(allTaskIDs) == 0 {
+			return delegateMessages
+		}
+		currentTasks, err := getAgentMessageTaskRows(allTaskIDs)
+		if err != nil {
+			logging.LogError(err, "Failed to fetch delegated tasking")
+			return delegateMessages
+		}
+		tasksByCallbackID := make(map[int][]agentMessageGetTaskingTask)
+		issuedTaskIDs := make([]int, 0, len(currentTasks))
+		for _, currentTask := range currentTasks {
+			if _, ok := routablePaths[currentTask.CallbackID]; !ok {
+				continue
+			}
+			tasksByCallbackID[currentTask.CallbackID] = append(tasksByCallbackID[currentTask.CallbackID], buildAgentMessageTask(currentTask))
+			issuedTaskIDs = append(issuedTaskIDs, currentTask.ID)
+		}
+		if len(issuedTaskIDs) == 0 {
+			return delegateMessages
+		}
+		if err := markAgentMessageTasksProcessing(issuedTaskIDs, time.Now().UTC()); err != nil {
+			logging.LogError(err, "Failed to update delegated task status to processing")
+			return delegateMessages
+		}
+		submittedTasksAwaitingFetching.removeTasksAfterProcessingUpdate(issuedTaskIDs)
+		for _, taskID := range issuedTaskIDs {
+			go addMitreAttackTaskMapping(taskID)
+		}
+		for _, targetCallbackId := range routableCallbackIds {
+			tasks := tasksByCallbackID[targetCallbackId]
+			if len(tasks) == 0 {
+				continue
+			}
+			routablePath := routablePaths[targetCallbackId]
+			newTask := map[string]interface{}{
+				"action": "get_tasking",
+				"tasks":  tasks,
+			}
+			wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime)
+			if err != nil {
+				logging.LogError(err, "Failed to recursively encrypt message")
+				continue
+			}
+			delegateMessages = append(delegateMessages, delegateMessageResponse{
+				Message:       string(wrappedMessage),
+				SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
+				C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
+			})
+		}
 	}
 	return delegateMessages
 }
@@ -208,59 +276,41 @@ func getDelegateProxyMessages(callbackID int, agentUUIDLength int, updateCheckin
 		for _, targetCallbackId := range callbackIds {
 			if routablePath := callbackGraph.GetBFSPath(callbackID, targetCallbackId); routablePath != nil && len(routablePath) > 0 {
 				// there's a route between our callback and the target callback for some sort of proxy data
-				if messages, err := proxyPorts.GetDataForCallbackId(targetCallbackId, CALLBACK_PORT_TYPE_SOCKS); err != nil {
-					logging.LogError(err, "Failed to get socks proxy data for routable callback")
-				} else if messages != nil {
-					// now that we have a path, need to recursively encrypt and wrap
-					newTask := map[string]interface{}{
-						"action":                 "get_tasking",
-						CALLBACK_PORT_TYPE_SOCKS: messages,
-					}
-					if wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime); err != nil {
-						logging.LogError(err, "Failed to recursively encrypt message")
-					} else {
-						delegateMessages = append(delegateMessages, delegateMessageResponse{
-							Message:       string(wrappedMessage),
-							SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
-							C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
-						})
-					}
+				proxyData, err := proxyPorts.GetDataForCallbackIdAllTypes(targetCallbackId)
+				if err != nil {
+					logging.LogError(err, "Failed to get proxy data for routable callback")
+					continue
 				}
-				if messages, err := proxyPorts.GetDataForCallbackId(targetCallbackId, CALLBACK_PORT_TYPE_RPORTFWD); err != nil {
-					logging.LogError(err, "Failed to get rpfwd proxy data for routable callback")
-				} else if messages != nil {
-					// now that we have a path, need to recursively encrypt and wrap
-					newTask := map[string]interface{}{
-						"action":                    "get_tasking",
-						CALLBACK_PORT_TYPE_RPORTFWD: messages,
-					}
-					if wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime); err != nil {
-						logging.LogError(err, "Failed to recursively encrypt message")
-					} else {
-						delegateMessages = append(delegateMessages, delegateMessageResponse{
-							Message:       string(wrappedMessage),
-							SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
-							C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
-						})
-					}
+				newTask := map[string]interface{}{
+					"action": "get_tasking",
 				}
-				if messages, err := proxyPorts.GetDataForCallbackId(targetCallbackId, CALLBACK_PORT_TYPE_INTERACTIVE); err != nil {
-					logging.LogError(err, "Failed to get interactive proxy data for routable callback")
-				} else if messages != nil {
+				gotData := false
+				if len(proxyData.Socks) > 0 {
 					// now that we have a path, need to recursively encrypt and wrap
-					newTask := map[string]interface{}{
-						"action":                       "get_tasking",
-						CALLBACK_PORT_TYPE_INTERACTIVE: messages,
-					}
-					if wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime); err != nil {
+					newTask[CALLBACK_PORT_TYPE_SOCKS] = proxyData.Socks
+					gotData = true
+				}
+				if len(proxyData.Rpfwd) > 0 {
+					// now that we have a path, need to recursively encrypt and wrap
+					newTask[CALLBACK_PORT_TYPE_RPORTFWD] = proxyData.Rpfwd
+					gotData = true
+				}
+				if len(proxyData.Interactive) > 0 {
+					// now that we have a path, need to recursively encrypt and wrap
+					newTask[CALLBACK_PORT_TYPE_INTERACTIVE] = proxyData.Interactive
+					gotData = true
+				}
+				if gotData {
+					wrappedMessage, err := RecursivelyEncryptMessage(routablePath, newTask, updateCheckinTime)
+					if err != nil {
 						logging.LogError(err, "Failed to recursively encrypt message")
-					} else {
-						delegateMessages = append(delegateMessages, delegateMessageResponse{
-							Message:       string(wrappedMessage),
-							SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
-							C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
-						})
+						continue
 					}
+					delegateMessages = append(delegateMessages, delegateMessageResponse{
+						Message:       string(wrappedMessage),
+						SuppliedUuid:  routablePath[len(routablePath)-1].DestinationAgentId,
+						C2ProfileName: routablePath[len(routablePath)-1].C2ProfileName,
+					})
 				}
 			}
 		}
@@ -7,7 +7,6 @@ import (
 	"github.com/its-a-feature/Mythic/database"
 	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
 	"github.com/its-a-feature/Mythic/logging"
-	"github.com/its-a-feature/Mythic/utils"
 	"github.com/jmoiron/sqlx"
 	"slices"
 	"sync"
@@ -19,6 +18,7 @@ import (
 type cbGraph struct {
 	// map callback_id -> array of adjacent callback_ids
 	adjMatrix map[int][]cbGraphAdjMatrixEntry
+	version   uint64
 	lock      sync.RWMutex
 }
 type cbGraphAdjMatrixEntry struct {
@@ -48,52 +48,54 @@ var uuidToIdAndOpId map[string]idAndOpId

 type bfsCache struct {
 	// source -> destination -> path
-	cache map[int]map[int][][]cbGraphAdjMatrixEntry
+	cache map[int]map[int]bfsCacheEntry
 	lock  sync.RWMutex
 }

+type bfsCacheEntry struct {
+	graphVersion uint64
+	path         []cbGraphAdjMatrixEntry
+}
+
 var BFSCache bfsCache

-func (c *bfsCache) GetPath(sourceId int, destinationId int) []cbGraphAdjMatrixEntry {
+func (c *bfsCache) GetPath(sourceId int, destinationId int, graphVersion uint64) ([]cbGraphAdjMatrixEntry, bool) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
 	if _, sourceExists := c.cache[sourceId]; sourceExists {
-		if _, destinationExists := c.cache[sourceId][destinationId]; destinationExists {
-			if len(c.cache[sourceId][destinationId]) > 0 {
-				//logging.LogInfo("Got path from cache", "source", sourceId, "destination", destinationId, "path", c.cache[sourceId][destinationId][0])
-				return c.cache[sourceId][destinationId][0]
+		if cachedPath, destinationExists := c.cache[sourceId][destinationId]; destinationExists {
+			if cachedPath.graphVersion == graphVersion {
+				return cachedPath.path, true
 			}
-			return nil
-		} else {
-			return nil
 		}
-	} else {
-		return nil
 	}
+	return nil, false
 }
-func (c *bfsCache) Remove(sourceId int) {
+func (c *bfsCache) Reset() {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-	if _, sourceExists := c.cache[sourceId]; sourceExists {
-		delete(c.cache, sourceId)
-	}
-	return
+	c.cache = make(map[int]map[int]bfsCacheEntry)
 }
-func (c *bfsCache) Add(sourceId int, destinationId int, bfsPath []cbGraphAdjMatrixEntry) {
+func (c *bfsCache) Add(sourceId int, destinationId int, graphVersion uint64, bfsPath []cbGraphAdjMatrixEntry) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
+	if c.cache == nil {
+		c.cache = make(map[int]map[int]bfsCacheEntry)
+	}
 	if _, sourceExists := c.cache[sourceId]; !sourceExists {
-		c.cache[sourceId] = make(map[int][][]cbGraphAdjMatrixEntry)
+		c.cache[sourceId] = make(map[int]bfsCacheEntry)
 	}
-	if _, destinationExists := c.cache[sourceId][destinationId]; !destinationExists {
-		c.cache[sourceId][destinationId] = [][]cbGraphAdjMatrixEntry{
-			bfsPath,
-		}
-	} else {
-		c.cache[sourceId][destinationId] = append(c.cache[sourceId][destinationId], bfsPath)
+	c.cache[sourceId][destinationId] = bfsCacheEntry{
+		graphVersion: graphVersion,
+		path:         append([]cbGraphAdjMatrixEntry(nil), bfsPath...),
 	}
 }

+func (g *cbGraph) bumpVersionLocked() {
+	g.version += 1
+	BFSCache.Reset()
+}
+
 func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 	g.lock.RLock()
 	defer g.lock.RUnlock()
@@ -211,7 +213,7 @@ func (g *cbGraph) Initialize() {
 			g.Add(edge.Destination, edge.Source, edge.C2Profile.Name, true)
 		}
 	}
-	BFSCache.cache = make(map[int]map[int][][]cbGraphAdjMatrixEntry)
+	BFSCache.Reset()
 }
 func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStructs.Callback, c2profileName string, initializing bool) {
 	c2 := getC2ProfileForName(c2profileName)
@@ -219,6 +221,9 @@ func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStruc
 		return
 	}
 	g.lock.Lock()
+	if g.adjMatrix == nil {
+		g.adjMatrix = make(map[int][]cbGraphAdjMatrixEntry)
+	}
 	if _, ok := g.adjMatrix[source.ID]; !ok {
 		// add it
 		//logging.LogInfo("adding new adjMatrix connection", "source", source.ID, "destination", destination.ID, "c2 profile", c2profileName)
@@ -229,6 +234,7 @@ func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStruc
 			SourceAgentId:      source.AgentCallbackID,
 			C2ProfileName:      c2profileName,
 		})
+		g.bumpVersionLocked()
 		g.lock.Unlock()
 		return
 	}
@@ -259,6 +265,7 @@ func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStruc
 		SourceAgentId:      source.AgentCallbackID,
 		C2ProfileName:      c2profileName,
 	})
+	g.bumpVersionLocked()

 	g.lock.Unlock()
 	logging.LogInfo("adding new adjMatrix connection", "source", source.ID, "destination", destination.ID, "c2 profile", c2profileName, "adj", g.adjMatrix[source.ID])
@@ -362,11 +369,13 @@ func (g *cbGraph) Remove(sourceId int, destinationId int, c2profileName string)
 			//g.adjMatrix[sourceId][len(g.adjMatrix[sourceId])-1] = cbGraphAdjMatrixEntry{}
 			//g.adjMatrix[sourceId] = g.adjMatrix[sourceId][:len(g.adjMatrix[sourceId])-1]
 			logging.LogDebug("removed adj matrix entry", "sourceID", sourceId, "destinationID", destinationId, "new adj", g.adjMatrix[sourceId])
-			BFSCache.Remove(sourceId)
+			g.bumpVersionLocked()
 		}
 	}
 }
 func (g *cbGraph) CanHaveDelegates(sourceId int) bool {
+	g.lock.RLock()
+	defer g.lock.RUnlock()
 	immediateChildren, exists := g.adjMatrix[sourceId]
 	if !exists {
 		return false
@@ -375,21 +384,21 @@ func (g *cbGraph) CanHaveDelegates(sourceId int) bool {
 }
 func (g *cbGraph) GetBFSPath(sourceId int, destinationId int) []cbGraphAdjMatrixEntry {
 	// breadth-first search from g.adjMatrix[sourceId] ->
-	var visitedIDs []int
 	var finalPath []cbGraphAdjMatrixEntry
 	if sourceId == destinationId {
 		return nil
 	}
-	if existingBFSpath := BFSCache.GetPath(sourceId, destinationId); existingBFSpath != nil {
+	g.lock.RLock()
+	defer g.lock.RUnlock()
+	graphVersion := g.version
+	if existingBFSpath, cached := BFSCache.GetPath(sourceId, destinationId, graphVersion); cached {
 		//logging.LogDebug("returning a cachedBFS path", "sourceID", sourceId, "destinationID", destinationId)
 		return existingBFSpath
 	}
 	// start with a fake-ish entry that we won't include in the final path
 	options := []cbGraphBFSEntry{{ID: sourceId, ParentBFS: nil}}
-	g.lock.RLock()
-	defer g.lock.RUnlock()
+	queuedIDs := map[int]bool{sourceId: true}
 	for index := 0; index < len(options); index++ {
-		visitedIDs = append(visitedIDs, options[index].ID)
 		if options[index].ID == destinationId {
 			// we found our entry, now follow the path back up for the shortest path
 			parent := options[index].ParentBFS
@@ -399,7 +408,7 @@ func (g *cbGraph) GetBFSPath(sourceId int, destinationId int) []cbGraphAdjMatrix
 				finalPath = append(finalPath, parent.MatrixEntry)
 				parent = parent.ParentBFS
 			}
-			BFSCache.Add(sourceId, destinationId, finalPath)
+			BFSCache.Add(sourceId, destinationId, graphVersion, finalPath)
 			return finalPath
 		}
 		// look through cbGraph.adjMatrix[options[index].ID] to see if we have our destination
@@ -409,18 +418,20 @@ func (g *cbGraph) GetBFSPath(sourceId int, destinationId int) []cbGraphAdjMatrix
 			continue
 		} else {
 			for _, nextChild := range nextChildren {
-				if !utils.SliceContains(visitedIDs, nextChild.DestinationId) {
-					options = append(options, cbGraphBFSEntry{
-						ID:          nextChild.DestinationId,
-						ParentBFS:   &options[index],
-						MatrixEntry: nextChild,
-					})
+				if queuedIDs[nextChild.DestinationId] {
+					continue
 				}
+				queuedIDs[nextChild.DestinationId] = true
+				options = append(options, cbGraphBFSEntry{
+					ID:          nextChild.DestinationId,
+					ParentBFS:   &options[index],
+					MatrixEntry: nextChild,
+				})
 			}
 		}
 	}
 	// add a known "empty" so we don't keep looking
-	BFSCache.Add(sourceId, destinationId, []cbGraphAdjMatrixEntry{})
+	BFSCache.Add(sourceId, destinationId, graphVersion, []cbGraphAdjMatrixEntry{})
 	return []cbGraphAdjMatrixEntry{}
 }
 func (g *cbGraph) Print() {
@@ -0,0 +1,149 @@
+package rabbitmq
+
+import (
+	"fmt"
+	"testing"
+
+	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
+)
+
+const testC2ProfileName = "test-c2"
+
+func newTestCallbackGraph(t *testing.T) *cbGraph {
+	t.Helper()
+	previousC2ProfileNameToIDMap := c2profileNameToIdMap
+	BFSCache.Reset()
+	c2profileNameToIdMap = map[string]databaseStructs.C2profile{
+		testC2ProfileName: {
+			ID:    1,
+			Name:  testC2ProfileName,
+			IsP2p: false,
+		},
+	}
+	t.Cleanup(func() {
+		BFSCache.Reset()
+		c2profileNameToIdMap = previousC2ProfileNameToIDMap
+	})
+	return &cbGraph{
+		adjMatrix: make(map[int][]cbGraphAdjMatrixEntry),
+	}
+}
+
+func testCallback(id int) databaseStructs.Callback {
+	return databaseStructs.Callback{
+		ID:              id,
+		AgentCallbackID: fmt.Sprintf("callback-%d", id),
+	}
+}
+
+func addTestGraphEdge(graph *cbGraph, sourceID int, destinationID int) {
+	graph.Add(testCallback(sourceID), testCallback(destinationID), testC2ProfileName, true)
+}
+
+func assertBFSPath(t *testing.T, path []cbGraphAdjMatrixEntry, expectedEdges [][2]int) {
+	t.Helper()
+	if len(path) != len(expectedEdges) {
+		t.Fatalf("expected path length %d, got %d: %#v", len(expectedEdges), len(path), path)
+	}
+	for i, expectedEdge := range expectedEdges {
+		if path[i].SourceId != expectedEdge[0] || path[i].DestinationId != expectedEdge[1] {
+			t.Fatalf("expected path[%d] to be %d -> %d, got %d -> %d",
+				i, expectedEdge[0], expectedEdge[1], path[i].SourceId, path[i].DestinationId)
+		}
+	}
+}
+
+func TestGetBFSPathReturnsShortestPathInDelegateWrappingOrder(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 2)
+	addTestGraphEdge(graph, 2, 4)
+	addTestGraphEdge(graph, 1, 3)
+	addTestGraphEdge(graph, 3, 5)
+	addTestGraphEdge(graph, 5, 4)
+
+	path := graph.GetBFSPath(1, 4)
+
+	assertBFSPath(t, path, [][2]int{
+		{2, 4},
+		{1, 2},
+	})
+}
+
+func TestGetBFSPathInvalidatesCachedNoPathWhenEdgeIsAdded(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 2)
+
+	path := graph.GetBFSPath(1, 3)
+	if len(path) != 0 {
+		t.Fatalf("expected no path before adding edge, got %#v", path)
+	}
+	if cachedPath, cached := BFSCache.GetPath(1, 3, graph.version); !cached || len(cachedPath) != 0 {
+		t.Fatalf("expected cached no-path result, cached=%v path=%#v", cached, cachedPath)
+	}
+
+	addTestGraphEdge(graph, 2, 3)
+	path = graph.GetBFSPath(1, 3)
+
+	assertBFSPath(t, path, [][2]int{
+		{2, 3},
+		{1, 2},
+	})
+}
+
+func TestGetBFSPathInvalidatesCachedPathWhenIntermediateEdgeIsRemoved(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 2)
+	addTestGraphEdge(graph, 2, 3)
+
+	path := graph.GetBFSPath(1, 3)
+	assertBFSPath(t, path, [][2]int{
+		{2, 3},
+		{1, 2},
+	})
+	versionBeforeRemove := graph.version
+
+	graph.Remove(2, 3, testC2ProfileName)
+	if graph.version == versionBeforeRemove {
+		t.Fatalf("expected graph version to change after removing an edge")
+	}
+	path = graph.GetBFSPath(1, 3)
+	if len(path) != 0 {
+		t.Fatalf("expected no path after removing intermediate edge, got %#v", path)
+	}
+}
+
+func TestGetBFSPathRefreshesCachedLongerPathWhenShorterEdgeIsAdded(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 2)
+	addTestGraphEdge(graph, 2, 4)
+
+	path := graph.GetBFSPath(1, 4)
+	assertBFSPath(t, path, [][2]int{
+		{2, 4},
+		{1, 2},
+	})
+
+	addTestGraphEdge(graph, 1, 4)
+	path = graph.GetBFSPath(1, 4)
+
+	assertBFSPath(t, path, [][2]int{
+		{1, 4},
+	})
+}
+
+func TestCanHaveDelegatesUsesCurrentGraphState(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	if graph.CanHaveDelegates(1) {
+		t.Fatalf("expected callback without edges to have no delegates")
+	}
+
+	addTestGraphEdge(graph, 1, 2)
+	if !graph.CanHaveDelegates(1) {
+		t.Fatalf("expected callback with an outbound edge to have delegates")
+	}
+
+	graph.Remove(1, 2, testC2ProfileName)
+	if graph.CanHaveDelegates(1) {
+		t.Fatalf("expected callback with removed edge to have no delegates")
+	}
+}
@@ -68,7 +68,12 @@ type submittedTask struct {
 	OperationID         int
 }
 type submittedTasksForAgents struct {
-	Tasks *[]submittedTask
+	Tasks                         *[]submittedTask
+	tasksByCallbackID             map[int][]int
+	interactiveTasksByCallbackID  map[int][]int
+	callbacksWithTasks            []int
+	callbacksWithInteractiveTasks []int
+	taskByID                      map[int]submittedTask
 	sync.RWMutex
 }

@@ -77,7 +82,10 @@ var submittedTasksAwaitingFetching submittedTasksForAgents
 func (s *submittedTasksForAgents) Initialize() {
 	currentTasks := []databaseStructs.Task{}
 	taskArray := make([]submittedTask, 0)
+	s.Lock()
 	s.Tasks = &taskArray
+	s.resetTaskIndexesLocked()
+	s.Unlock()
 	if err := database.DB.Select(&currentTasks, `SELECT
 	id, callback_id, operation_id, is_interactive_task, interactive_task_type, parent_task_id
 	FROM task
@@ -86,18 +94,125 @@ func (s *submittedTasksForAgents) Initialize() {
 	} else {
 		s.Lock()
 		for _, t := range currentTasks {
-			*s.Tasks = append(*s.Tasks, submittedTask{
-				TaskID:              t.ID,
-				CallbackID:          t.CallbackID,
-				OperationID:         t.OperationID,
-				IsInteractiveTask:   t.IsInteractiveTask,
-				InteractiveTaskType: int(t.InteractiveTaskType.Int64),
-				ParentTaskID:        int(t.ParentTaskID.Int64),
-			})
+			s.addTaskLocked(newSubmittedTask(t))
 		}
 		s.Unlock()
 	}
 }
+
+func newSubmittedTask(task databaseStructs.Task) submittedTask {
+	return submittedTask{
+		TaskID:              task.ID,
+		CallbackID:          task.CallbackID,
+		OperationID:         task.OperationID,
+		IsInteractiveTask:   task.IsInteractiveTask,
+		InteractiveTaskType: int(task.InteractiveTaskType.Int64),
+		ParentTaskID:        int(task.ParentTaskID.Int64),
+	}
+}
+
+func (s *submittedTasksForAgents) resetTaskIndexesLocked() {
+	s.tasksByCallbackID = make(map[int][]int)
+	s.interactiveTasksByCallbackID = make(map[int][]int)
+	s.callbacksWithTasks = make([]int, 0)
+	s.callbacksWithInteractiveTasks = make([]int, 0)
+	s.taskByID = make(map[int]submittedTask)
+}
+
+func (s *submittedTasksForAgents) ensureTaskStoreLocked() {
+	if s.Tasks == nil {
+		taskArray := make([]submittedTask, 0)
+		s.Tasks = &taskArray
+	}
+	if s.taskByID == nil || s.tasksByCallbackID == nil || s.interactiveTasksByCallbackID == nil {
+		tasks := *s.Tasks
+		s.resetTaskIndexesLocked()
+		for _, task := range tasks {
+			s.addTaskToIndexesLocked(task)
+		}
+	}
+}
+
+func (s *submittedTasksForAgents) addTaskLocked(task submittedTask) bool {
+	s.ensureTaskStoreLocked()
+	if _, exists := s.taskByID[task.TaskID]; exists {
+		return false
+	}
+	*s.Tasks = append(*s.Tasks, task)
+	s.addTaskToIndexesLocked(task)
+	return true
+}
+
+func (s *submittedTasksForAgents) addTaskToIndexesLocked(task submittedTask) {
+	s.taskByID[task.TaskID] = task
+	if task.IsInteractiveTask {
+		if len(s.interactiveTasksByCallbackID[task.CallbackID]) == 0 {
+			s.callbacksWithInteractiveTasks = append(s.callbacksWithInteractiveTasks, task.CallbackID)
+		}
+		s.interactiveTasksByCallbackID[task.CallbackID] = append(s.interactiveTasksByCallbackID[task.CallbackID], task.TaskID)
+	} else {
+		if len(s.tasksByCallbackID[task.CallbackID]) == 0 {
+			s.callbacksWithTasks = append(s.callbacksWithTasks, task.CallbackID)
+		}
+		s.tasksByCallbackID[task.CallbackID] = append(s.tasksByCallbackID[task.CallbackID], task.TaskID)
+	}
+}
+
+func (s *submittedTasksForAgents) removeTaskFromIndexesLocked(task submittedTask) {
+	delete(s.taskByID, task.TaskID)
+	if task.IsInteractiveTask {
+		s.interactiveTasksByCallbackID[task.CallbackID] = removeTaskID(s.interactiveTasksByCallbackID[task.CallbackID], task.TaskID)
+		if len(s.interactiveTasksByCallbackID[task.CallbackID]) == 0 {
+			delete(s.interactiveTasksByCallbackID, task.CallbackID)
+			s.callbacksWithInteractiveTasks = removeCallbackID(s.callbacksWithInteractiveTasks, task.CallbackID)
+		}
+		return
+	}
+	s.tasksByCallbackID[task.CallbackID] = removeTaskID(s.tasksByCallbackID[task.CallbackID], task.TaskID)
+	if len(s.tasksByCallbackID[task.CallbackID]) == 0 {
+		delete(s.tasksByCallbackID, task.CallbackID)
+		s.callbacksWithTasks = removeCallbackID(s.callbacksWithTasks, task.CallbackID)
+	}
+}
+
+func removeTaskID(taskIDs []int, taskID int) []int {
+	for i, currentTaskID := range taskIDs {
+		if currentTaskID == taskID {
+			return append(taskIDs[:i], taskIDs[i+1:]...)
+		}
+	}
+	return taskIDs
+}
+
+func removeCallbackID(callbackIDs []int, callbackID int) []int {
+	for i, currentCallbackID := range callbackIDs {
+		if currentCallbackID == callbackID {
+			return append(callbackIDs[:i], callbackIDs[i+1:]...)
+		}
+	}
+	return callbackIDs
+}
+
+func cloneTaskIDs(taskIDs []int) []int {
+	return append([]int(nil), taskIDs...)
+}
+
+func (s *submittedTasksForAgents) removeTaskFromStoreLocked(taskID int) (submittedTask, bool) {
+	s.ensureTaskStoreLocked()
+	task, ok := s.taskByID[taskID]
+	if !ok {
+		return submittedTask{}, false
+	}
+	s.removeTaskFromIndexesLocked(task)
+	for i := 0; i < len(*s.Tasks); i++ {
+		if (*s.Tasks)[i].TaskID == taskID {
+			*s.Tasks = append((*s.Tasks)[:i], (*s.Tasks)[i+1:]...)
+			break
+		}
+	}
+	return task, true
+}
+
 func (s *submittedTasksForAgents) addTask(task databaseStructs.Task) {
 	// before we add this for agents to pick up, see if this task can go to a pushC2 connected callback
 	if grpc.PushC2Server.CheckClientConnected(task.CallbackID) > PushC2Connections.Connected {
@@ -169,14 +284,7 @@ func (s *submittedTasksForAgents) addTask(task databaseStructs.Task) {
 	s.Lock()
 	defer s.Unlock()
 	logging.LogInfo("adding task to SubmittedTasks")
-	*s.Tasks = append(*s.Tasks, submittedTask{
-		TaskID:              task.ID,
-		CallbackID:          task.CallbackID,
-		OperationID:         task.OperationID,
-		IsInteractiveTask:   task.IsInteractiveTask,
-		InteractiveTaskType: int(task.InteractiveTaskType.Int64),
-		ParentTaskID:        int(task.ParentTaskID.Int64),
-	})
+	s.addTaskLocked(newSubmittedTask(task))
 }
 func (s *submittedTasksForAgents) addTaskById(taskId int) {
 	task := databaseStructs.Task{ID: taskId}
@@ -190,66 +298,67 @@ func (s *submittedTasksForAgents) addTaskById(taskId int) {
 }
 func (s *submittedTasksForAgents) removeTask(taskId int) {
 	s.Lock()
-	defer s.Unlock()
-	for i := 0; i < len(*s.Tasks); i++ {
-		if (*s.Tasks)[i].TaskID == taskId {
-			EventingChannel <- EventNotification{
-				Trigger:     eventing.TriggerTaskStart,
-				OperationID: (*s.Tasks)[i].OperationID,
-				TaskID:      (*s.Tasks)[i].TaskID,
-			}
-			if (*s.Tasks)[i].IsInteractiveTask {
-				_, err := database.DB.Exec(`UPDATE task SET
+	task, found := s.removeTaskFromStoreLocked(taskId)
+	s.Unlock()
+	if !found {
+		return
+	}
+	EventingChannel <- EventNotification{
+		Trigger:     eventing.TriggerTaskStart,
+		OperationID: task.OperationID,
+		TaskID:      task.TaskID,
+	}
+	if task.IsInteractiveTask {
+		_, err := database.DB.Exec(`UPDATE task SET
                status_timestamp_processing=$1, status_timestamp_processed=$1, completed=true 
                WHERE id=$2`, time.Now().UTC(), taskId)
-				if err != nil {
-					logging.LogError(err, "failed to update timestamp for interactive task")
-				}
-			} else {
-				_, err := database.DB.Exec(`UPDATE task SET
+		if err != nil {
+			logging.LogError(err, "failed to update timestamp for interactive task")
+		}
+	} else {
+		_, err := database.DB.Exec(`UPDATE task SET
                status_timestamp_processing=$1 
                WHERE id=$2`, time.Now().UTC(), taskId)
-				if err != nil {
-					logging.LogError(err, "failed to update timestamp for interactive task")
-				}
-			}
-
-			*s.Tasks = append((*s.Tasks)[:i], (*s.Tasks)[i+1:]...)
-			return
+		if err != nil {
+			logging.LogError(err, "failed to update timestamp for interactive task")
 		}
 	}
 }
 func (s *submittedTasksForAgents) getTasksForCallbackId(callbackId int) []int {
-	tasks := []int{}
 	s.RLock()
 	defer s.RUnlock()
-	for i := 0; i < len(*s.Tasks); i++ {
-		if (*s.Tasks)[i].CallbackID == callbackId && !(*s.Tasks)[i].IsInteractiveTask {
-			tasks = append(tasks, (*s.Tasks)[i].TaskID)
+	return cloneTaskIDs(s.tasksByCallbackID[callbackId])
+}
+func (s *submittedTasksForAgents) getTasksForCallbackIds(callbackIds []int) map[int][]int {
+	tasksByCallbackID := make(map[int][]int)
+	if len(callbackIds) == 0 {
+		return tasksByCallbackID
+	}
+	callbackIDMap := make(map[int]bool, len(callbackIds))
+	for _, callbackId := range callbackIds {
+		callbackIDMap[callbackId] = true
+	}
+	s.RLock()
+	defer s.RUnlock()
+	for callbackID := range callbackIDMap {
+		if taskIDs := s.tasksByCallbackID[callbackID]; len(taskIDs) > 0 {
+			tasksByCallbackID[callbackID] = cloneTaskIDs(taskIDs)
 		}
 	}
-	return tasks
+	return tasksByCallbackID
 }
 func (s *submittedTasksForAgents) getInteractiveTasksForCallbackId(callbackId int) []int {
-	tasks := []int{}
 	s.RLock()
 	defer s.RUnlock()
-	for i := 0; i < len(*s.Tasks); i++ {
-		if (*s.Tasks)[i].CallbackID == callbackId && (*s.Tasks)[i].IsInteractiveTask {
-			tasks = append(tasks, (*s.Tasks)[i].TaskID)
-		}
-	}
-	return tasks
+	return cloneTaskIDs(s.interactiveTasksByCallbackID[callbackId])
 }
 func (s *submittedTasksForAgents) getOtherCallbackIds(callbackId int) []int {
 	callbacks := []int{}
 	s.RLock()
 	defer s.RUnlock()
-	for i := 0; i < len(*s.Tasks); i++ {
-		if (*s.Tasks)[i].CallbackID != callbackId && !(*s.Tasks)[i].IsInteractiveTask {
-			if !utils.SliceContains(callbacks, (*s.Tasks)[i].CallbackID) {
-				callbacks = append(callbacks, (*s.Tasks)[i].CallbackID)
-			}
+	for _, currentCallbackID := range s.callbacksWithTasks {
+		if currentCallbackID != callbackId {
+			callbacks = append(callbacks, currentCallbackID)
 		}
 	}
 	return callbacks
@@ -258,15 +367,44 @@ func (s *submittedTasksForAgents) getInteractiveTasksOtherCallbackIds(callbackId
 	callbacks := []int{}
 	s.RLock()
 	defer s.RUnlock()
-	for i := 0; i < len(*s.Tasks); i++ {
-		if (*s.Tasks)[i].CallbackID != callbackId && (*s.Tasks)[i].IsInteractiveTask {
-			if !utils.SliceContains(callbacks, (*s.Tasks)[i].CallbackID) {
-				callbacks = append(callbacks, (*s.Tasks)[i].CallbackID)
-			}
+	for _, currentCallbackID := range s.callbacksWithInteractiveTasks {
+		if currentCallbackID != callbackId {
+			callbacks = append(callbacks, currentCallbackID)
 		}
 	}
 	return callbacks
 }
+func (s *submittedTasksForAgents) removeTasksAfterProcessingUpdate(taskIds []int) {
+	if len(taskIds) == 0 {
+		return
+	}
+	taskIDMap := make(map[int]bool, len(taskIds))
+	for _, taskId := range taskIds {
+		taskIDMap[taskId] = true
+	}
+	taskStartEvents := make([]EventNotification, 0, len(taskIds))
+	s.Lock()
+	s.ensureTaskStoreLocked()
+	pendingTasks := (*s.Tasks)[:0]
+	for i := 0; i < len(*s.Tasks); i++ {
+		task := (*s.Tasks)[i]
+		if taskIDMap[task.TaskID] {
+			taskStartEvents = append(taskStartEvents, EventNotification{
+				Trigger:     eventing.TriggerTaskStart,
+				OperationID: task.OperationID,
+				TaskID:      task.TaskID,
+			})
+			s.removeTaskFromIndexesLocked(task)
+			continue
+		}
+		pendingTasks = append(pendingTasks, task)
+	}
+	*s.Tasks = pendingTasks
+	s.Unlock()
+	for _, taskStartEvent := range taskStartEvents {
+		EventingChannel <- taskStartEvent
+	}
+}

 func CreateTask(createTaskInput CreateTaskInput) CreateTaskResponse {
 	response := CreateTaskResponse{
@@ -0,0 +1,87 @@
+package rabbitmq
+
+import (
+	"slices"
+	"testing"
+)
+
+func drainEventingChannelForTest() {
+	for {
+		select {
+		case <-EventingChannel:
+		default:
+			return
+		}
+	}
+}
+
+func addSubmittedTaskForTest(taskStore *submittedTasksForAgents, task submittedTask) {
+	taskStore.Lock()
+	defer taskStore.Unlock()
+	taskStore.addTaskLocked(task)
+}
+
+func TestSubmittedTasksForAgentsUsesCallbackIndexes(t *testing.T) {
+	taskStore := submittedTasksForAgents{}
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 1, CallbackID: 10, OperationID: 100})
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 2, CallbackID: 10, OperationID: 100, IsInteractiveTask: true})
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 3, CallbackID: 20, OperationID: 100})
+
+	if taskIDs := taskStore.getTasksForCallbackId(10); !slices.Equal(taskIDs, []int{1}) {
+		t.Fatalf("expected callback 10 non-interactive task IDs [1], got %#v", taskIDs)
+	}
+	if taskIDs := taskStore.getInteractiveTasksForCallbackId(10); !slices.Equal(taskIDs, []int{2}) {
+		t.Fatalf("expected callback 10 interactive task IDs [2], got %#v", taskIDs)
+	}
+	if callbackIDs := taskStore.getOtherCallbackIds(10); !slices.Equal(callbackIDs, []int{20}) {
+		t.Fatalf("expected other callback IDs [20], got %#v", callbackIDs)
+	}
+
+	tasksByCallbackID := taskStore.getTasksForCallbackIds([]int{10, 20})
+	if !slices.Equal(tasksByCallbackID[10], []int{1}) {
+		t.Fatalf("expected callback 10 task IDs [1], got %#v", tasksByCallbackID[10])
+	}
+	if !slices.Equal(tasksByCallbackID[20], []int{3}) {
+		t.Fatalf("expected callback 20 task IDs [3], got %#v", tasksByCallbackID[20])
+	}
+
+	taskIDs := taskStore.getTasksForCallbackId(10)
+	taskIDs[0] = 99
+	if taskIDs = taskStore.getTasksForCallbackId(10); !slices.Equal(taskIDs, []int{1}) {
+		t.Fatalf("expected returned task ID slice to be a copy, got %#v", taskIDs)
+	}
+}
+
+func TestSubmittedTasksForAgentsBatchRemoveUpdatesIndexes(t *testing.T) {
+	drainEventingChannelForTest()
+	t.Cleanup(drainEventingChannelForTest)
+
+	taskStore := submittedTasksForAgents{}
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 1, CallbackID: 10, OperationID: 100})
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 2, CallbackID: 10, OperationID: 100})
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 3, CallbackID: 10, OperationID: 100, IsInteractiveTask: true})
+	addSubmittedTaskForTest(&taskStore, submittedTask{TaskID: 4, CallbackID: 20, OperationID: 200})
+
+	taskStore.removeTasksAfterProcessingUpdate([]int{1, 4})
+
+	if taskIDs := taskStore.getTasksForCallbackId(10); !slices.Equal(taskIDs, []int{2}) {
+		t.Fatalf("expected callback 10 non-interactive task IDs [2], got %#v", taskIDs)
+	}
+	if taskIDs := taskStore.getInteractiveTasksForCallbackId(10); !slices.Equal(taskIDs, []int{3}) {
+		t.Fatalf("expected callback 10 interactive task IDs [3], got %#v", taskIDs)
+	}
+	if callbackIDs := taskStore.getOtherCallbackIds(10); len(callbackIDs) != 0 {
+		t.Fatalf("expected no other non-interactive callback IDs, got %#v", callbackIDs)
+	}
+
+	taskStore.RLock()
+	if len(*taskStore.Tasks) != 2 {
+		t.Fatalf("expected 2 tasks left in backing store, got %d", len(*taskStore.Tasks))
+	}
+	taskStore.RUnlock()
+
+	taskStartEvents := []EventNotification{<-EventingChannel, <-EventingChannel}
+	if taskStartEvents[0].TaskID != 1 || taskStartEvents[1].TaskID != 4 {
+		t.Fatalf("expected task start events for task IDs 1 and 4, got %#v", taskStartEvents)
+	}
+}
@@ -12,6 +12,7 @@ import (
 	"net"
 	"sort"
 	"strconv"
+	"sync"
 	"time"

 	"github.com/its-a-feature/Mythic/database"
@@ -91,9 +92,19 @@ type callbackPortUsage struct {

 type callbackPortsInUse struct {
 	ports                        []*callbackPortUsage
+	portsByCallbackID            map[int][]*callbackPortUsage
+	portsByCallbackIDAndType     map[int]map[CallbackPortType][]*callbackPortUsage
+	callbacksWithPorts           []int
 	proxyFromAgentMessageChannel chan ProxyFromAgentMessageForMythic
 	bytesReceivedFromAgentChan   chan bytesReceivedFromAgentMessage
 	bytesSentToAgentChan         chan bytesSentToAgentMessage
+	sync.RWMutex
+}
+
+type callbackProxyData struct {
+	Socks       []proxyToAgentMessage
+	Rpfwd       []proxyToAgentMessage
+	Interactive []agentMessagePostResponseInteractive
 }

 var proxyPorts callbackPortsInUse
@@ -212,10 +223,13 @@ func ManuallyToggleProxy(input ProxyStop) ProxyStopResponse {
 }
 func (c *callbackPortsInUse) Initialize() {
 	callbackPorts := []databaseStructs.Callbackport{}
+	c.Lock()
 	c.ports = make([]*callbackPortUsage, 0)
+	c.resetPortIndexesLocked()
 	c.proxyFromAgentMessageChannel = make(chan ProxyFromAgentMessageForMythic, 2000)
 	c.bytesReceivedFromAgentChan = make(chan bytesReceivedFromAgentMessage, 2000)
 	c.bytesSentToAgentChan = make(chan bytesSentToAgentMessage, 2000)
+	c.Unlock()
 	go c.ListenForProxyFromAgentMessage()
 	go c.ListenForNewByteTransferUpdates()
 	if err := database.DB.Select(&callbackPorts, `SELECT 
@@ -267,11 +281,155 @@ func (c *callbackPortsInUse) Initialize() {
 			if err := newPort.Start(); err != nil {
 				logging.LogError(err, "Failed to start listening", "port info", &newPort)
 			} else {
-				c.ports = append(c.ports, &newPort)
+				c.Lock()
+				c.addPortLocked(&newPort)
+				c.Unlock()
 			}
 		}
 	}
 }
+
+func (c *callbackPortsInUse) resetPortIndexesLocked() {
+	c.portsByCallbackID = make(map[int][]*callbackPortUsage)
+	c.portsByCallbackIDAndType = make(map[int]map[CallbackPortType][]*callbackPortUsage)
+	c.callbacksWithPorts = make([]int, 0)
+}
+
+func (c *callbackPortsInUse) ensurePortIndexesLocked() {
+	if c.portsByCallbackID == nil || c.portsByCallbackIDAndType == nil {
+		ports := c.ports
+		c.resetPortIndexesLocked()
+		for _, port := range ports {
+			c.addPortToIndexesLocked(port)
+		}
+	}
+}
+
+func (c *callbackPortsInUse) addPortLocked(port *callbackPortUsage) {
+	c.ensurePortIndexesLocked()
+	c.ports = append(c.ports, port)
+	c.addPortToIndexesLocked(port)
+}
+
+func (c *callbackPortsInUse) addPortToIndexesLocked(port *callbackPortUsage) {
+	if len(c.portsByCallbackID[port.CallbackID]) == 0 {
+		c.callbacksWithPorts = append(c.callbacksWithPorts, port.CallbackID)
+	}
+	c.portsByCallbackID[port.CallbackID] = append(c.portsByCallbackID[port.CallbackID], port)
+	if _, ok := c.portsByCallbackIDAndType[port.CallbackID]; !ok {
+		c.portsByCallbackIDAndType[port.CallbackID] = make(map[CallbackPortType][]*callbackPortUsage)
+	}
+	c.portsByCallbackIDAndType[port.CallbackID][port.PortType] = append(c.portsByCallbackIDAndType[port.CallbackID][port.PortType], port)
+}
+
+func (c *callbackPortsInUse) removePortLocked(port *callbackPortUsage) bool {
+	c.ensurePortIndexesLocked()
+	found := false
+	for i := 0; i < len(c.ports); i++ {
+		if c.ports[i] == port {
+			c.ports = append(c.ports[:i], c.ports[i+1:]...)
+			found = true
+			break
+		}
+	}
+	if found {
+		c.removePortFromIndexesLocked(port)
+	}
+	return found
+}
+
+func (c *callbackPortsInUse) removePortFromIndexesLocked(port *callbackPortUsage) {
+	c.portsByCallbackID[port.CallbackID] = removeCallbackPortUsage(c.portsByCallbackID[port.CallbackID], port)
+	if len(c.portsByCallbackID[port.CallbackID]) == 0 {
+		delete(c.portsByCallbackID, port.CallbackID)
+		delete(c.portsByCallbackIDAndType, port.CallbackID)
+		c.callbacksWithPorts = removeProxyCallbackID(c.callbacksWithPorts, port.CallbackID)
+		return
+	}
+	c.portsByCallbackIDAndType[port.CallbackID][port.PortType] = removeCallbackPortUsage(c.portsByCallbackIDAndType[port.CallbackID][port.PortType], port)
+	if len(c.portsByCallbackIDAndType[port.CallbackID][port.PortType]) == 0 {
+		delete(c.portsByCallbackIDAndType[port.CallbackID], port.PortType)
+	}
+}
+
+func removeCallbackPortUsage(ports []*callbackPortUsage, port *callbackPortUsage) []*callbackPortUsage {
+	for i, currentPort := range ports {
+		if currentPort == port {
+			return append(ports[:i], ports[i+1:]...)
+		}
+	}
+	return ports
+}
+
+func removeProxyCallbackID(callbackIDs []int, callbackID int) []int {
+	for i, currentCallbackID := range callbackIDs {
+		if currentCallbackID == callbackID {
+			return append(callbackIDs[:i], callbackIDs[i+1:]...)
+		}
+	}
+	return callbackIDs
+}
+
+func cloneCallbackPortUsages(ports []*callbackPortUsage) []*callbackPortUsage {
+	return append([]*callbackPortUsage(nil), ports...)
+}
+
+func (c *callbackPortsInUse) getPortsForCallbackAndType(callbackID int, portType CallbackPortType) []*callbackPortUsage {
+	c.RLock()
+	defer c.RUnlock()
+	if c.portsByCallbackIDAndType == nil {
+		return c.getPortsForCallbackAndTypeFromSliceLocked(callbackID, portType)
+	}
+	return cloneCallbackPortUsages(c.portsByCallbackIDAndType[callbackID][portType])
+}
+
+func (c *callbackPortsInUse) getPortsForCallbackAndTypeFromSliceLocked(callbackID int, portType CallbackPortType) []*callbackPortUsage {
+	ports := make([]*callbackPortUsage, 0)
+	for _, port := range c.ports {
+		if port.CallbackID == callbackID && port.PortType == portType {
+			ports = append(ports, port)
+		}
+	}
+	return ports
+}
+
+func (c *callbackPortsInUse) getPortsForCallbackByType(callbackID int, portTypes []CallbackPortType) map[CallbackPortType][]*callbackPortUsage {
+	portsByType := make(map[CallbackPortType][]*callbackPortUsage, len(portTypes))
+	c.RLock()
+	defer c.RUnlock()
+	if c.portsByCallbackIDAndType == nil {
+		for _, portType := range portTypes {
+			portsByType[portType] = c.getPortsForCallbackAndTypeFromSliceLocked(callbackID, portType)
+		}
+		return portsByType
+	}
+	for _, portType := range portTypes {
+		portsByType[portType] = cloneCallbackPortUsages(c.portsByCallbackIDAndType[callbackID][portType])
+	}
+	return portsByType
+}
+
+func (c *callbackPortsInUse) findPortForRemoval(callbackId int, portType CallbackPortType, localPort int, operationId int, remoteIP string, remotePort int,
+	username string, password string) *callbackPortUsage {
+	c.RLock()
+	defer c.RUnlock()
+	ports := c.portsByCallbackIDAndType[callbackId][portType]
+	if c.portsByCallbackIDAndType == nil {
+		ports = c.getPortsForCallbackAndTypeFromSliceLocked(callbackId, portType)
+	}
+	for _, port := range ports {
+		if port.OperationID == operationId &&
+			port.Username == username &&
+			port.Password == password &&
+			port.LocalPort == localPort {
+			if remoteIP != "" && remotePort != 0 && (remoteIP != port.RemoteIP || remotePort != port.RemotePort) {
+				continue
+			}
+			return port
+		}
+	}
+	return nil
+}
 func (c *callbackPortsInUse) ListenForNewByteTransferUpdates() {
 	oldByteValues := make(map[int]map[string]int64)
 	currentByteValues := make(map[int]map[string]int64)
@@ -358,40 +516,34 @@ func (c *callbackPortsInUse) ListenForProxyFromAgentMessage() {
 		case CALLBACK_PORT_TYPE_RPORTFWD:
 			fallthrough
 		case CALLBACK_PORT_TYPE_SOCKS:
+			ports := c.getPortsForCallbackAndType(agentMessage.CallbackID, agentMessage.PortType)
 			// loop through each message and find the corresponding callback + local port combo
 			for j := 0; j < len(agentMessage.Messages); j++ {
-				for i := 0; i < len(c.ports); i++ {
-					if c.ports[i].CallbackID == agentMessage.CallbackID && c.ports[i].PortType == agentMessage.PortType {
-						if agentMessage.Messages[j].Port > 0 {
-							// port is specified, try to find the right one
-							if c.ports[i].LocalPort == agentMessage.Messages[j].Port {
-								c.ports[i].messagesFromAgent <- agentMessage.Messages[j]
-								break
-							}
-						} else {
-							// didn't specify a specific port, to just send it to the first matching type
-							c.ports[i].messagesFromAgent <- agentMessage.Messages[j]
+				for _, port := range ports {
+					if agentMessage.Messages[j].Port > 0 {
+						// port is specified, try to find the right one
+						if port.LocalPort == agentMessage.Messages[j].Port {
+							port.messagesFromAgent <- agentMessage.Messages[j]
 							break
 						}
+					} else {
+						// didn't specify a specific port, to just send it to the first matching type
+						port.messagesFromAgent <- agentMessage.Messages[j]
+						break
 					}
 				}
 				//logging.LogInfo("got message from agent", "chan", messages[j].ServerID, "p.messagesFromAgentQueue", len(c.ports[i].messagesFromAgent))
 				//c.ports[i].messagesFromAgent <- agentMessage.Messages[j]
 			}
 		case CALLBACK_PORT_TYPE_INTERACTIVE:
-			foundPort := false
-			for i := 0; i < len(c.ports); i++ {
-				if c.ports[i].CallbackID == agentMessage.CallbackID && c.ports[i].PortType == agentMessage.PortType {
-					for j := 0; j < len(agentMessage.InteractiveMessages); j++ {
-						//logging.LogInfo("got message from agent", "chan", messages[j].ServerID, "p.messagesFromAgentQueue", len(c.ports[i].messagesFromAgent))
-						c.ports[i].interactiveMessagesFromAgent <- agentMessage.InteractiveMessages[j]
-					}
-					handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
-					foundPort = true
-					break
+			ports := c.getPortsForCallbackAndType(agentMessage.CallbackID, agentMessage.PortType)
+			if len(ports) > 0 {
+				for j := 0; j < len(agentMessage.InteractiveMessages); j++ {
+					//logging.LogInfo("got message from agent", "chan", messages[j].ServerID, "p.messagesFromAgentQueue", len(c.ports[i].messagesFromAgent))
+					ports[0].interactiveMessagesFromAgent <- agentMessage.InteractiveMessages[j]
 				}
-			}
-			if !foundPort {
+				handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
+			} else {
 				go handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
 			}

@@ -399,7 +551,7 @@ func (c *callbackPortsInUse) ListenForProxyFromAgentMessage() {
 	}
 }
 func (c *callbackPortsInUse) GetNextAvailableLocalPort() uint32 {
-	exposedPorts := utils.MythicConfig.ServerDynamicPorts
+	exposedPorts := append([]uint32(nil), utils.MythicConfig.ServerDynamicPorts...)
 	sort.Slice(exposedPorts, func(i, j int) bool {
 		return exposedPorts[i] < exposedPorts[j]
 	})
@@ -411,15 +563,17 @@ func (c *callbackPortsInUse) GetNextAvailableLocalPort() uint32 {
 	return 0
 }
 func (c *callbackPortsInUse) IsPortInUse(port uint32) bool {
-	for i, _ := range c.ports {
-		switch c.ports[i].PortType {
+	c.RLock()
+	defer c.RUnlock()
+	for _, currentPort := range c.ports {
+		switch currentPort.PortType {
 		case CALLBACK_PORT_TYPE_RPORTFWD:
 			// rpfwd doesn't bind a port on Mythic, so the ports don't count here
 			continue
 		case CALLBACK_PORT_TYPE_SOCKS:
 			fallthrough
 		case CALLBACK_PORT_TYPE_INTERACTIVE:
-			if uint32(c.ports[i].LocalPort) == port {
+			if uint32(currentPort.LocalPort) == port {
 				return true
 			}
 		}
@@ -427,17 +581,16 @@ func (c *callbackPortsInUse) IsPortInUse(port uint32) bool {
 	return false
 }
 func (c *callbackPortsInUse) GetPortForTypeAndCallback(taskId int, callbackId int, portType CallbackPortType) int {
-	for i, _ := range c.ports {
+	ports := c.getPortsForCallbackAndType(callbackId, portType)
+	for _, port := range ports {
 		switch portType {
 		case CALLBACK_PORT_TYPE_SOCKS:
-			if c.ports[i].PortType == portType && c.ports[i].CallbackID == callbackId {
-				return c.ports[i].LocalPort
-			}
+			return port.LocalPort
 		case CALLBACK_PORT_TYPE_RPORTFWD:
 			fallthrough
 		case CALLBACK_PORT_TYPE_INTERACTIVE:
-			if c.ports[i].PortType == portType && c.ports[i].CallbackID == callbackId && c.ports[i].TaskID == taskId {
-				return c.ports[i].LocalPort
+			if port.TaskID == taskId {
+				return port.LocalPort
 			}
 		}

@@ -445,46 +598,79 @@ func (c *callbackPortsInUse) GetPortForTypeAndCallback(taskId int, callbackId in
 	return 0
 }
 func (c *callbackPortsInUse) GetDataForCallbackIdPortType(callbackId int, portType CallbackPortType) (interface{}, error) {
-	var interactiveData []agentMessagePostResponseInteractive
-	var socksData []proxyToAgentMessage
-	fetchedInteractive := false
-	for i := 0; i < len(c.ports); i++ {
-		if c.ports[i].CallbackID == callbackId && c.ports[i].PortType == portType {
-			switch portType {
-			case CALLBACK_PORT_TYPE_INTERACTIVE:
-				interactiveData = append(interactiveData, c.ports[i].GetData().([]agentMessagePostResponseInteractive)...)
-				fetchedInteractive = true
-			case CALLBACK_PORT_TYPE_SOCKS:
-				fallthrough
-			case CALLBACK_PORT_TYPE_RPORTFWD:
-				socksData = append(socksData, c.ports[i].GetData().([]proxyToAgentMessage)...)
-			}
-
-		}
-	}
-	if portType == CALLBACK_PORT_TYPE_INTERACTIVE && !fetchedInteractive {
-		newInteractiveData, err := handleAgentMessageGetInteractiveTasking(callbackId)
-		if err != nil {
-			logging.LogError(err, "Failed to fetch interactive tasks")
-		} else {
-			interactiveData = append(interactiveData, newInteractiveData...)
-		}
+	proxyData, err := c.getDataForCallbackIdPortTypes(callbackId, []CallbackPortType{portType})
+	if err != nil {
+		return nil, err
 	}
 	switch portType {
 	case CALLBACK_PORT_TYPE_INTERACTIVE:
-		if len(interactiveData) > 0 {
-			return interactiveData, nil
+		if len(proxyData.Interactive) > 0 {
+			return proxyData.Interactive, nil
 		}
 	case CALLBACK_PORT_TYPE_SOCKS:
-		fallthrough
+		if len(proxyData.Socks) > 0 {
+			return proxyData.Socks, nil
+		}
 	case CALLBACK_PORT_TYPE_RPORTFWD:
-		if len(socksData) > 0 {
-			return socksData, nil
+		if len(proxyData.Rpfwd) > 0 {
+			return proxyData.Rpfwd, nil
 		}
 	}

 	return nil, nil
 }
+
+func (c *callbackPortsInUse) GetDataForCallbackIdAllTypes(callbackId int) (callbackProxyData, error) {
+	return c.getDataForCallbackIdPortTypes(callbackId, []CallbackPortType{
+		CALLBACK_PORT_TYPE_SOCKS,
+		CALLBACK_PORT_TYPE_RPORTFWD,
+		CALLBACK_PORT_TYPE_INTERACTIVE,
+	})
+}
+
+func (c *callbackPortsInUse) getDataForCallbackIdPortTypes(callbackId int, portTypes []CallbackPortType) (callbackProxyData, error) {
+	proxyData := callbackProxyData{}
+	portsByType := c.getPortsForCallbackByType(callbackId, portTypes)
+	for _, portType := range portTypes {
+		ports := portsByType[portType]
+		switch portType {
+		case CALLBACK_PORT_TYPE_INTERACTIVE:
+			for _, port := range ports {
+				proxyData.Interactive = append(proxyData.Interactive, port.GetData().([]agentMessagePostResponseInteractive)...)
+			}
+			if len(ports) == 0 {
+				newInteractiveData, err := handleAgentMessageGetInteractiveTasking(callbackId)
+				if err != nil {
+					logging.LogError(err, "Failed to fetch interactive tasks")
+				} else {
+					proxyData.Interactive = append(proxyData.Interactive, newInteractiveData...)
+				}
+			}
+		case CALLBACK_PORT_TYPE_SOCKS:
+			for _, port := range ports {
+				proxyData.Socks = append(proxyData.Socks, port.GetData().([]proxyToAgentMessage)...)
+			}
+		case CALLBACK_PORT_TYPE_RPORTFWD:
+			for _, port := range ports {
+				proxyData.Rpfwd = append(proxyData.Rpfwd, port.GetData().([]proxyToAgentMessage)...)
+			}
+		}
+	}
+	return proxyData, nil
+}
+
+func (d callbackProxyData) AddToResponse(response map[string]interface{}) {
+	if len(d.Socks) > 0 {
+		response[CALLBACK_PORT_TYPE_SOCKS] = d.Socks
+	}
+	if len(d.Rpfwd) > 0 {
+		response[CALLBACK_PORT_TYPE_RPORTFWD] = d.Rpfwd
+	}
+	if len(d.Interactive) > 0 {
+		response[CALLBACK_PORT_TYPE_INTERACTIVE] = d.Interactive
+	}
+}
+
 func (c *callbackPortsInUse) SendDataToCallbackIdPortType(callbackId int, portType CallbackPortType, messages []proxyFromAgentMessage) {
 	c.proxyFromAgentMessageChannel <- ProxyFromAgentMessageForMythic{
 		CallbackID: callbackId,
@@ -501,9 +687,21 @@ func (c *callbackPortsInUse) SendInteractiveDataToCallbackIdPortType(callbackId
 }
 func (c *callbackPortsInUse) GetOtherCallbackIds(callbackId int) []int {
 	callbackIds := []int{}
-	for i := 0; i < len(c.ports); i++ {
-		if c.ports[i].CallbackID != callbackId {
-			callbackIds = append(callbackIds, c.ports[i].CallbackID)
+	c.RLock()
+	defer c.RUnlock()
+	if c.portsByCallbackID == nil {
+		callbackIDMap := make(map[int]bool)
+		for _, port := range c.ports {
+			if port.CallbackID != callbackId && !callbackIDMap[port.CallbackID] {
+				callbackIDMap[port.CallbackID] = true
+				callbackIds = append(callbackIds, port.CallbackID)
+			}
+		}
+		return callbackIds
+	}
+	for _, currentCallbackID := range c.callbacksWithPorts {
+		if currentCallbackID != callbackId {
+			callbackIds = append(callbackIds, currentCallbackID)
 		}
 	}
 	return callbackIds
@@ -595,7 +793,9 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 		}
 		return err
 	}
-	c.ports = append(c.ports, &newPort)
+	c.Lock()
+	c.addPortLocked(&newPort)
+	c.Unlock()
 	return nil

 }
@@ -620,37 +820,25 @@ func (c *callbackPortsInUse) Test(callbackId int, portType CallbackPortType, loc
 }
 func (c *callbackPortsInUse) Remove(callbackId int, portType CallbackPortType, localPort int, operationId int, remoteIP string, remotePort int,
 	username string, password string) error {
-	for i := 0; i < len(c.ports); i++ {
-		if c.ports[i].CallbackID == callbackId &&
-			c.ports[i].OperationID == operationId &&
-			c.ports[i].PortType == portType &&
-			c.ports[i].Username == username &&
-			c.ports[i].Password == password &&
-			c.ports[i].LocalPort == localPort {
-			if remoteIP != "" && remotePort != 0 && (remoteIP != c.ports[i].RemoteIP || remotePort != c.ports[i].RemotePort) {
-				continue
-			}
-			err := c.ports[i].Stop()
-			if err != nil {
-				logging.LogError(err, "Failed to stop proxy")
-				//c.Unlock()
-				return err
-			}
-			queryString := `UPDATE callbackport SET deleted=true WHERE id=$1`
-			queryArgs := []interface{}{c.ports[i].CallbackPortID}
-			_, err = database.DB.Exec(queryString, queryArgs...)
-			if err != nil {
-				logging.LogError(err, "Failed to delete port mapping from database for proxy")
-				//c.Unlock()
-				return err
-			}
-			c.ports = append((c.ports)[:i], (c.ports)[i+1:]...)
-			//c.Unlock()
-			return nil
-
-		}
+	port := c.findPortForRemoval(callbackId, portType, localPort, operationId, remoteIP, remotePort, username, password)
+	if port == nil {
+		return nil
 	}
-	//c.Unlock()
+	err := port.Stop()
+	if err != nil {
+		logging.LogError(err, "Failed to stop proxy")
+		return err
+	}
+	queryString := `UPDATE callbackport SET deleted=true WHERE id=$1`
+	queryArgs := []interface{}{port.CallbackPortID}
+	_, err = database.DB.Exec(queryString, queryArgs...)
+	if err != nil {
+		logging.LogError(err, "Failed to delete port mapping from database for proxy")
+		return err
+	}
+	c.Lock()
+	c.removePortLocked(port)
+	c.Unlock()
 	return nil
 }
 func isPortExposedThroughDocker(portToCheck int) bool {
@@ -0,0 +1,94 @@
+package rabbitmq
+
+import (
+	"slices"
+	"testing"
+)
+
+func newTestCallbackPort(callbackID int, portType CallbackPortType, localPort int, taskID int) *callbackPortUsage {
+	return &callbackPortUsage{
+		CallbackID:                   callbackID,
+		PortType:                     portType,
+		LocalPort:                    localPort,
+		TaskID:                       taskID,
+		OperationID:                  1,
+		messagesToAgent:              make(chan proxyToAgentMessage, 10),
+		interactiveMessagesToAgent:   make(chan agentMessagePostResponseInteractive, 10),
+		messagesFromAgent:            make(chan proxyFromAgentMessage, 10),
+		interactiveMessagesFromAgent: make(chan agentMessagePostResponseInteractive, 10),
+		stopAllConnections:           make(chan bool, 1),
+	}
+}
+
+func addTestCallbackPort(registry *callbackPortsInUse, port *callbackPortUsage) {
+	registry.Lock()
+	defer registry.Unlock()
+	registry.addPortLocked(port)
+}
+
+func removeTestCallbackPort(registry *callbackPortsInUse, port *callbackPortUsage) {
+	registry.Lock()
+	defer registry.Unlock()
+	registry.removePortLocked(port)
+}
+
+func TestCallbackPortsInUseIndexesCallbacksAndPortTypes(t *testing.T) {
+	registry := callbackPortsInUse{}
+	socksPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	rpfwdPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_RPORTFWD, 7002, 2)
+	otherPort := newTestCallbackPort(20, CALLBACK_PORT_TYPE_SOCKS, 8001, 3)
+	addTestCallbackPort(&registry, socksPort)
+	addTestCallbackPort(&registry, rpfwdPort)
+	addTestCallbackPort(&registry, otherPort)
+
+	if callbackIDs := registry.GetOtherCallbackIds(10); !slices.Equal(callbackIDs, []int{20}) {
+		t.Fatalf("expected callback IDs [20], got %#v", callbackIDs)
+	}
+	if port := registry.GetPortForTypeAndCallback(0, 10, CALLBACK_PORT_TYPE_SOCKS); port != 7001 {
+		t.Fatalf("expected socks local port 7001, got %d", port)
+	}
+	if port := registry.GetPortForTypeAndCallback(2, 10, CALLBACK_PORT_TYPE_RPORTFWD); port != 7002 {
+		t.Fatalf("expected rpfwd local port 7002, got %d", port)
+	}
+
+	removeTestCallbackPort(&registry, otherPort)
+	if callbackIDs := registry.GetOtherCallbackIds(10); len(callbackIDs) != 0 {
+		t.Fatalf("expected no other callback IDs after removing other port, got %#v", callbackIDs)
+	}
+}
+
+func TestCallbackPortsInUseCombinedDataFetchUsesIndexedPorts(t *testing.T) {
+	registry := callbackPortsInUse{}
+	socksPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	rpfwdPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_RPORTFWD, 7002, 2)
+	interactivePort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_INTERACTIVE, 7003, 3)
+	addTestCallbackPort(&registry, socksPort)
+	addTestCallbackPort(&registry, rpfwdPort)
+	addTestCallbackPort(&registry, interactivePort)
+
+	socksPort.messagesToAgent <- proxyToAgentMessage{ServerID: 1, Message: []byte("socks")}
+	rpfwdPort.messagesToAgent <- proxyToAgentMessage{ServerID: 2, Message: []byte("rpfwd")}
+	interactivePort.interactiveMessagesToAgent <- agentMessagePostResponseInteractive{TaskUUID: "interactive"}
+
+	proxyData, err := registry.GetDataForCallbackIdAllTypes(10)
+	if err != nil {
+		t.Fatalf("unexpected error fetching proxy data: %v", err)
+	}
+	if len(proxyData.Socks) != 1 || string(proxyData.Socks[0].Message) != "socks" {
+		t.Fatalf("expected one socks message, got %#v", proxyData.Socks)
+	}
+	if len(proxyData.Rpfwd) != 1 || string(proxyData.Rpfwd[0].Message) != "rpfwd" {
+		t.Fatalf("expected one rpfwd message, got %#v", proxyData.Rpfwd)
+	}
+	if len(proxyData.Interactive) != 1 || proxyData.Interactive[0].TaskUUID != "interactive" {
+		t.Fatalf("expected one interactive message, got %#v", proxyData.Interactive)
+	}
+
+	proxyData, err = registry.GetDataForCallbackIdAllTypes(10)
+	if err != nil {
+		t.Fatalf("unexpected error fetching proxy data second time: %v", err)
+	}
+	if len(proxyData.Socks) != 0 || len(proxyData.Rpfwd) != 0 || len(proxyData.Interactive) != 0 {
+		t.Fatalf("expected proxy data to be drained, got %#v", proxyData)
+	}
+}