From 751999a82b8854ef5231e0ded5fe022b1ea5d8ff Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sat, 4 Jul 2026 18:44:17 +0900 Subject: Harden backup import and restore coverage --- src/cli/export.c | 48 ++- src/cli/import.c | 580 ++++++++++++++++++++++++------------ src/domain/search.c | 20 +- src/domain/search.h | 1 + src/mcp/tools/common/history.c | 19 ++ src/mcp/tools/common/history.h | 1 + src/mcp/tools/common/memory_index.c | 19 ++ src/mcp/tools/common/memory_index.h | 1 + src/mcp/tools/get_system_health.c | 19 ++ src/storage/store.c | 23 ++ src/storage/store.h | 3 + 11 files changed, 542 insertions(+), 192 deletions(-) (limited to 'src') diff --git a/src/cli/export.c b/src/cli/export.c index 4288ba4..d7397c4 100644 --- a/src/cli/export.c +++ b/src/cli/export.c @@ -3,17 +3,42 @@ #include "mcp/tools/common/common.h" #include "storage/codec.h" #include "storage/store.h" +#include "sys/system.h" #include #include #include +static bool +export_finish(FILE *output, const char *temporary_path, const char *destination) +{ + bool ok = fflush(output) == 0; + if (output == stdout) + { + return ok; + } + if (fclose(output) != 0) + { + ok = false; + } + if (ok && !tiny_sys_path_replace(temporary_path, destination)) + { + ok = false; + } + if (!ok) + { + (void)tiny_sys_file_remove(temporary_path); + } + return ok; +} + int tiny_cli_export(tiny_store_t *store, const tiny_cli_options_t *options) { size_t index; size_t emitted = 0; FILE *output = NULL; + char *temporary_path = NULL; tiny_json_buffer_t item; if (!options || !options->export_path) @@ -28,9 +53,19 @@ tiny_cli_export(tiny_store_t *store, const tiny_cli_options_t *options) } else { - if (fopen_s(&output, options->export_path, "wb") != 0 || !output) + size_t path_length = strlen(options->export_path); + temporary_path = malloc(path_length + sizeof(".tmp")); + if (!temporary_path) + { + fprintf(stderr, "export: out of memory\n"); + return 1; + } + memcpy(temporary_path, options->export_path, path_length); + memcpy(temporary_path + path_length, ".tmp", sizeof(".tmp")); + if (fopen_s(&output, temporary_path, "wb") != 0 || !output) { fprintf(stderr, "export: unable to open '%s'\n", options->export_path); + free(temporary_path); return 1; } } @@ -41,7 +76,9 @@ tiny_cli_export(tiny_store_t *store, const tiny_cli_options_t *options) if (output != stdout) { fclose(output); + (void)tiny_sys_file_remove(temporary_path); } + free(temporary_path); return 1; } @@ -83,16 +120,21 @@ tiny_cli_export(tiny_store_t *store, const tiny_cli_options_t *options) if (output != stdout) { fclose(output); + (void)tiny_sys_file_remove(temporary_path); } + free(temporary_path); return 1; } emitted++; } - if (output != stdout) + if (!export_finish(output, temporary_path, options->export_path)) { - fclose(output); + fprintf(stderr, "export: unable to complete '%s'\n", options->export_path); + free(temporary_path); + return 1; } + free(temporary_path); (void)emitted; return 0; } diff --git a/src/cli/import.c b/src/cli/import.c index 3ddb9e9..f4688ab 100644 --- a/src/cli/import.c +++ b/src/cli/import.c @@ -9,8 +9,11 @@ #include #include #include +#include #define IMPORT_MAX_LINE TINY_CLI_LINE_MAX +#define IMPORT_MAX_RECORDS 100000U +#define IMPORT_MAX_DECODED_BYTES (UINT64_C(4) * 1024U * 1024U * 1024U) typedef struct import_record_list { @@ -64,27 +67,264 @@ import_get_string(const char *json, int len, const char *path, char *out, size_t return decoded > 0 && (size_t)decoded < out_size; } +static int +import_hex_digit(char value) +{ + if (value >= '0' && value <= '9') + { + return value - '0'; + } + if (value >= 'a' && value <= 'f') + { + return value - 'a' + 10; + } + if (value >= 'A' && value <= 'F') + { + return value - 'A' + 10; + } + return -1; +} + +static bool +import_decode_json_string(const char *source, size_t length, char *output, size_t *output_length) +{ + size_t input = 0, used = 0; + while (input < length) + { + uint32_t codepoint; + unsigned count; + if (source[input] != '\\') + { + output[used++] = source[input++]; + continue; + } + if (++input >= length) + { + return false; + } + switch (source[input++]) + { + case '"': + output[used++] = '"'; + continue; + case '\\': + output[used++] = '\\'; + continue; + case '/': + output[used++] = '/'; + continue; + case 'b': + output[used++] = '\b'; + continue; + case 'f': + output[used++] = '\f'; + continue; + case 'n': + output[used++] = '\n'; + continue; + case 'r': + output[used++] = '\r'; + continue; + case 't': + output[used++] = '\t'; + continue; + case 'u': + break; + default: + return false; + } + if (input + 4 > length) + { + return false; + } + codepoint = 0; + for (count = 0; count < 4; ++count) + { + int digit = import_hex_digit(source[input++]); + if (digit < 0) + { + return false; + } + codepoint = (codepoint << 4) | (uint32_t)digit; + } + if (codepoint >= 0xd800 && codepoint <= 0xdbff) + { + uint32_t low = 0; + if (input + 6 > length || source[input] != '\\' || source[input + 1] != 'u') + { + return false; + } + input += 2; + for (count = 0; count < 4; ++count) + { + int digit = import_hex_digit(source[input++]); + if (digit < 0) + { + return false; + } + low = (low << 4) | (uint32_t)digit; + } + if (low < 0xdc00 || low > 0xdfff) + { + return false; + } + codepoint = UINT32_C(0x10000) + ((codepoint - 0xd800U) << 10) + low - 0xdc00U; + } + else if (codepoint >= 0xdc00 && codepoint <= 0xdfff) + { + return false; + } + if (codepoint <= 0x7f) + { + output[used++] = (char)codepoint; + } + else if (codepoint <= 0x7ff) + { + output[used++] = (char)(0xc0U | (codepoint >> 6)); + output[used++] = (char)(0x80U | (codepoint & 0x3fU)); + } + else if (codepoint <= 0xffff) + { + output[used++] = (char)(0xe0U | (codepoint >> 12)); + output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU)); + output[used++] = (char)(0x80U | (codepoint & 0x3fU)); + } + else + { + output[used++] = (char)(0xf0U | (codepoint >> 18)); + output[used++] = (char)(0x80U | ((codepoint >> 12) & 0x3fU)); + output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU)); + output[used++] = (char)(0x80U | (codepoint & 0x3fU)); + } + } + output[used] = '\0'; + *output_length = used; + return true; +} + +static bool +import_get_allocated_string(const char *json, + int len, + const char *path, + size_t max_length, + bool allow_empty, + char **out, + size_t *decoded_length) +{ + const char *token = NULL; + int token_length = 0; + size_t decoded; + char *value; + *out = NULL; + if (decoded_length) + { + *decoded_length = 0; + } + if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_STRING || token_length < 2 + || token_length == INT_MAX) + { + return false; + } + value = malloc((size_t)token_length); + if (!value) + { + return false; + } + if (!import_decode_json_string(token + 1, (size_t)token_length - 2U, value, &decoded) + || (!allow_empty && decoded == 0) || decoded > max_length || strlen(value) != decoded + || !tiny_utf8_valid(value, decoded)) + { + free(value); + return false; + } + *out = value; + if (decoded_length) + { + *decoded_length = decoded; + } + return true; +} + +static bool +import_get_integer_token(const char *json, + int len, + const char *path, + bool allow_negative, + uint64_t maximum, + uint64_t *magnitude, + bool *negative) +{ + const char *token = NULL; + int token_length = 0; + int index = 0; + uint64_t value = 0; + if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_NUMBER || token_length <= 0) + { + return false; + } + *negative = token[0] == '-'; + if (*negative) + { + if (!allow_negative || ++index == token_length) + { + return false; + } + } + if (token[index] == '0' && index + 1 != token_length) + { + return false; + } + for (; index < token_length; ++index) + { + unsigned digit; + if (token[index] < '0' || token[index] > '9') + { + return false; + } + digit = (unsigned)(token[index] - '0'); + if (value > (maximum - digit) / 10U) + { + return false; + } + value = value * 10U + digit; + } + *magnitude = value; + return true; +} + static bool import_get_i64(const char *json, int len, const char *path, int64_t *value) { - double number; - if (mjson_get_number(json, len, path, &number) != 1) + uint64_t magnitude; + bool negative; + if (!import_get_integer_token( + json, len, path, true, (uint64_t)INT64_MAX + 1U, &magnitude, &negative)) { return false; } - *value = (int64_t)number; + if (negative) + { + *value = magnitude == (uint64_t)INT64_MAX + 1U ? INT64_MIN : -(int64_t)magnitude; + } + else + { + if (magnitude > INT64_MAX) + { + return false; + } + *value = (int64_t)magnitude; + } return true; } static bool import_get_u64(const char *json, int len, const char *path, uint64_t *value) { - double number; - if (mjson_get_number(json, len, path, &number) != 1 || number < 0) + bool negative; + if (!import_get_integer_token(json, len, path, false, UINT64_MAX, value, &negative)) { return false; } - *value = (uint64_t)number; return true; } @@ -126,23 +366,15 @@ import_parse_string_array(const char *json, &value_type)) != 0) { - char value[512]; - int decoded; + char *value = NULL; + size_t decoded = 0; if (value_type != MJSON_TOK_STRING || index >= max_count - || (decoded = mjson_get_string( - array_json + value_offset, value_length, "$", value, sizeof(value))) - <= 0 - || (size_t)decoded > max_length || decoded >= (int)sizeof(value)) - { - break; - } - (*output)[index] = malloc((size_t)decoded + 1); - if (!(*output)[index]) + || !import_get_allocated_string( + array_json + value_offset, value_length, "$", max_length, false, &value, &decoded)) { break; } - memcpy((*output)[index], value, (size_t)decoded); - (*output)[index][decoded] = '\0'; + (*output)[index] = value; index++; } *count = index; @@ -183,31 +415,37 @@ import_parse_relations(const char *json, int len, const char *path, tiny_metadat &value_type)) != 0) { - char type_name[32]; - char message_id[32]; - int decoded; + char *type_name = NULL; + char *message_id = NULL; + size_t decoded = 0; tiny_relation_t *relation; if (value_type != MJSON_TOK_OBJECT || index >= TINY_RELATION_COUNT_MAX - || mjson_get_string( - array_json + value_offset, value_length, "$.type", type_name, sizeof(type_name)) - <= 0 - || (decoded = mjson_get_string(array_json + value_offset, - value_length, - "$.message_id", - message_id, - sizeof(message_id))) - <= 0) + || !import_get_allocated_string( + array_json + value_offset, value_length, "$.type", 31, false, &type_name, NULL) + || !import_get_allocated_string(array_json + value_offset, + value_length, + "$.message_id", + TINY_MESSAGE_ID_LENGTH, + false, + &message_id, + &decoded)) { + free(type_name); + free(message_id); return false; } relation = &metadata->relations[index]; if (!tiny_relation_type_parse(type_name, &relation->type) || !tiny_message_id_valid(message_id)) { + free(type_name); + free(message_id); return false; } memcpy(relation->message_id, message_id, (size_t)decoded); relation->message_id[decoded] = '\0'; + free(type_name); + free(message_id); index++; } metadata->relation_count = index; @@ -219,16 +457,30 @@ import_parse_metadata(const char *json, int len, tiny_message_t *message) { char tier_name[32] = { 0 }; char kind_name[32] = { 0 }; - char tool_call_id[TINY_TOOL_CALL_ID_MAX + 1] = { 0 }; - char summary[TINY_SUMMARY_MAX + 1] = { 0 }; - double importance_number = 0; + char *tool_call_id = NULL; + char *summary = NULL; + uint64_t importance_number = 0; + bool negative = false; bool has_tool_call_id, has_summary, has_importance; - has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0 - && import_get_string( - json, len, "$.metadata.tool_call_id", tool_call_id, sizeof(tool_call_id)); - has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0 - && import_get_string(json, len, "$.metadata.summary", summary, sizeof(summary)); - has_importance = mjson_get_number(json, len, "$.metadata.importance", &importance_number) == 1; + has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0; + has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0; + has_importance = mjson_find(json, len, "$.metadata.importance", NULL, NULL) != 0; + if ((has_tool_call_id + && !import_get_allocated_string( + json, len, "$.metadata.tool_call_id", TINY_TOOL_CALL_ID_MAX, false, &tool_call_id, NULL)) + || (has_summary + && !import_get_allocated_string( + json, len, "$.metadata.summary", TINY_SUMMARY_MAX, false, &summary, NULL)) + || (has_importance + && !import_get_integer_token( + json, len, "$.metadata.importance", false, 100, &importance_number, &negative))) + { + free(tool_call_id); + free(summary); + return false; + } + message->metadata.tool_call_id = tool_call_id; + message->metadata.summary = summary; if (!import_get_string(json, len, "$.metadata.memory_tier", tier_name, sizeof(tier_name)) || !import_get_string(json, len, "$.metadata.memory_kind", kind_name, sizeof(kind_name)) || !import_parse_string_array(json, @@ -254,30 +506,8 @@ import_parse_metadata(const char *json, int len, tiny_message_t *message) { return false; } - if (has_tool_call_id) - { - message->metadata.tool_call_id = malloc(strlen(tool_call_id) + 1); - if (!message->metadata.tool_call_id) - { - return false; - } - memcpy(message->metadata.tool_call_id, tool_call_id, strlen(tool_call_id) + 1); - } - if (has_summary) - { - message->metadata.summary = malloc(strlen(summary) + 1); - if (!message->metadata.summary) - { - return false; - } - memcpy(message->metadata.summary, summary, strlen(summary) + 1); - } if (has_importance) { - if (importance_number < 0 || importance_number > 100) - { - return false; - } message->metadata.importance = (uint8_t)importance_number; } return true; @@ -287,14 +517,14 @@ static bool import_manifest_valid(const char *line, int len) { char type_name[64]; - double version = 0; + uint64_t version = 0; if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name)) || strcmp(type_name, "tiny-convmem-manifest") != 0 - || mjson_get_number(line, len, "$.format_version", &version) != 1) + || !import_get_u64(line, len, "$.format_version", &version)) { return false; } - return (int)version == 2; + return version == 2; } static bool @@ -318,18 +548,21 @@ static bool import_parse_message_record(const char *line, int len, tiny_import_record_t *record) { char role_name[32]; - char session_id[TINY_SESSION_ID_MAX + 1]; - char content_buffer[4096]; + char *session_id = NULL; + char *content = NULL; + size_t content_length = 0; tiny_message_init(&record->message); record->deleted = false; if (!import_get_string( line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id)) - || !import_get_string(line, len, "$.session_id", session_id, sizeof(session_id)) || !import_get_string(line, len, "$.role", role_name, sizeof(role_name)) + || !import_get_allocated_string( + line, len, "$.session_id", TINY_SESSION_ID_MAX, false, &session_id, NULL) || !import_get_i64(line, len, "$.timestamp", &record->message.timestamp_ms) || !import_get_i64(line, len, "$.updated_at", &record->message.updated_at_ms) || !import_get_u64(line, len, "$.revision", &record->message.revision)) { + free(session_id); return false; } if (!tiny_message_id_valid(record->message.message_id) @@ -337,57 +570,18 @@ import_parse_message_record(const char *line, int len, tiny_import_record_t *rec || record->message.updated_at_ms < 0 || record->message.revision == 0 || record->message.revision > INT32_MAX) { + free(session_id); return false; } - if (mjson_find(line, len, "$.content", NULL, NULL) == 0) - { - return false; - } - if (mjson_get_string(line, len, "$.content", content_buffer, sizeof(content_buffer)) > 0) - { - size_t length = strlen(content_buffer); - if (length == 0 || length > TINY_CONTENT_MAX) - { - return false; - } - record->message.content = malloc(length + 1); - if (!record->message.content) - { - return false; - } - memcpy(record->message.content, content_buffer, length + 1); - } - else - { - const char *content = NULL; - int content_len = 0; - if (mjson_find(line, len, "$.content", &content, &content_len) == 0 || !content - || content_len < 2) - { - return false; - } - if ((size_t)(content_len - 2) == 0 || (size_t)(content_len - 2) > TINY_CONTENT_MAX) - { - return false; - } - record->message.content = malloc((size_t)content_len - 1); - if (!record->message.content) - { - return false; - } - memcpy(record->message.content, content + 1, (size_t)content_len - 2); - record->message.content[(size_t)content_len - 2] = '\0'; - } - if (!tiny_utf8_valid(record->message.content, strlen(record->message.content))) - { - return false; - } - record->message.session_id = malloc(strlen(session_id) + 1); - if (!record->message.session_id) + if (!import_get_allocated_string( + line, len, "$.content", TINY_CONTENT_MAX, false, &content, &content_length)) { + free(session_id); return false; } - memcpy(record->message.session_id, session_id, strlen(session_id) + 1); + (void)content_length; + record->message.content = content; + record->message.session_id = session_id; if (!import_parse_metadata(line, len, &record->message) || !tiny_metadata_valid(&record->message.metadata)) { @@ -398,60 +592,99 @@ import_parse_message_record(const char *line, int len, tiny_import_record_t *rec return true; } +static uint64_t +import_record_decoded_bytes(const tiny_import_record_t *record) +{ + uint64_t total = 0; + size_t index; + const tiny_message_t *message = &record->message; + if (record->deleted) + { + return TINY_MESSAGE_ID_LENGTH; + } + total = strlen(message->session_id) + strlen(message->content) + TINY_MESSAGE_ID_LENGTH; + total += message->metadata.tool_call_id ? strlen(message->metadata.tool_call_id) : 0; + total += message->metadata.summary ? strlen(message->metadata.summary) : 0; + for (index = 0; index < message->metadata.tag_count; ++index) + { + total += strlen(message->metadata.tags[index]); + } + total += message->metadata.relation_count * TINY_MESSAGE_ID_LENGTH; + total += message->metadata.source_count * TINY_MESSAGE_ID_LENGTH; + return total; +} + static bool -import_parse_jsonl(const char *jsonl, import_record_list_t *records) +import_parse_stream(FILE *file, import_record_list_t *records) { - const char *line_start = jsonl; - const char *line_end; + char *line = malloc((size_t)IMPORT_MAX_LINE + 2U); bool manifest_seen = false; - while (*line_start) + uint64_t decoded_total = 0; + if (!line) { - int length; + return false; + } + while (fgets(line, (int)IMPORT_MAX_LINE + 2, file) != NULL) + { + size_t length = strlen(line); tiny_import_record_t record; - line_end = strchr(line_start, '\n'); - if (line_end == NULL) + uint64_t decoded; + if (length == 0 || length > IMPORT_MAX_LINE + || (length == IMPORT_MAX_LINE + 1U && line[length - 1] != '\n') + || (line[length - 1] != '\n' && !feof(file))) { - line_end = line_start + strlen(line_start); + free(line); + return false; } - length = (int)(line_end - line_start); - if (length > 0) + if (length > 0 && line[length - 1] == '\n') { - if (!manifest_seen) - { - if (!import_manifest_valid(line_start, length)) - { - return false; - } - manifest_seen = true; - } - else if (import_parse_del_record(line_start, length, &record)) - { - if (!import_record_list_push(records)) - { - return false; - } - records->items[records->count++] = record; - } - else if (import_parse_message_record(line_start, length, &record)) - { - if (!import_record_list_push(records)) - { - tiny_message_free(&record.message); - return false; - } - records->items[records->count++] = record; - } - else + line[--length] = '\0'; + } + if (length > 0 && line[length - 1] == '\r') + { + line[--length] = '\0'; + } + if (length == 0) + { + continue; + } + if (!manifest_seen) + { + if (!import_manifest_valid(line, (int)length)) { + free(line); return false; } + manifest_seen = true; + continue; } - if (*line_end == '\0') + if (records->count >= IMPORT_MAX_RECORDS) { - break; + free(line); + return false; } - line_start = line_end + 1; + if (!import_parse_del_record(line, (int)length, &record) + && !import_parse_message_record(line, (int)length, &record)) + { + free(line); + return false; + } + decoded = import_record_decoded_bytes(&record); + if (decoded > IMPORT_MAX_DECODED_BYTES - decoded_total || !import_record_list_push(records)) + { + tiny_message_free(&record.message); + free(line); + return false; + } + decoded_total += decoded; + records->items[records->count++] = record; } + if (ferror(file)) + { + free(line); + return false; + } + free(line); return manifest_seen; } @@ -462,9 +695,7 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options) size_t imported = 0, rejected = 0; tiny_store_status_t status; tiny_import_policy_t policy = TINY_IMPORT_OVERWRITE; - char *jsonl = NULL; FILE *file = NULL; - long file_size; int result = 1; if (!options || !options->import_path) @@ -496,49 +727,22 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options) return 1; } - if (fseek(file, 0, SEEK_END) != 0 || (file_size = ftell(file)) < 0 - || fseek(file, 0, SEEK_SET) != 0) - { - fprintf(stderr, "import: unable to seek '%s'\n", options->import_path); - fclose(file); - return 1; - } - - if ((size_t)file_size > IMPORT_MAX_LINE * 1024ULL * 1024ULL) + if (!import_parse_stream(file, &records)) { - fprintf(stderr, "import: file too large\n"); - fclose(file); - return 1; - } - - jsonl = malloc((size_t)file_size + 1); - if (!jsonl) - { - fprintf(stderr, "import: out of memory\n"); - fclose(file); - return 1; - } - - if (fread(jsonl, 1, (size_t)file_size, file) != (size_t)file_size) - { - fprintf(stderr, "import: unable to read '%s'\n", options->import_path); - free(jsonl); + fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path); fclose(file); + import_record_list_free(&records); return 1; } - jsonl[file_size] = '\0'; - fclose(file); - - if (!import_parse_jsonl(jsonl, &records)) + if (fclose(file) != 0) { - fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path); - free(jsonl); + fprintf(stderr, "import: unable to close '%s'\n", options->import_path); + import_record_list_free(&records); return 1; } status = tiny_store_import(store, records.items, records.count, policy, &imported, &rejected); import_record_list_free(&records); - free(jsonl); if (status == TINY_STORE_OK) { diff --git a/src/domain/search.c b/src/domain/search.c index faf16ff..e34cc72 100644 --- a/src/domain/search.c +++ b/src/domain/search.c @@ -356,6 +356,24 @@ tiny_search_index_create(void) return calloc(1, sizeof(tiny_search_index_t)); } +uint64_t +tiny_search_index_memory_estimate(const tiny_search_index_t *index) +{ + uint64_t total; + size_t position; + if (!index) + { + return 0; + } + total = sizeof(*index) + index->bucket_capacity * sizeof(*index->buckets) + + index->slot_capacity * sizeof(*index->slots); + for (position = 0; position < index->bucket_count; ++position) + { + total += index->buckets[position].capacity * sizeof(*index->buckets[position].messages); + } + return total; +} + void tiny_search_index_free(tiny_search_index_t *index) { @@ -710,4 +728,4 @@ tiny_search_score(const char *query, const tiny_message_t *message) return 0.0; } return breakdown.total; -} \ No newline at end of file +} diff --git a/src/domain/search.h b/src/domain/search.h index a891a27..821102e 100644 --- a/src/domain/search.h +++ b/src/domain/search.h @@ -17,6 +17,7 @@ bool tiny_search_index_candidates(const tiny_search_index_t *index, const char *query, size_t **output, size_t *count); +uint64_t tiny_search_index_memory_estimate(const tiny_search_index_t *index); typedef struct tiny_search_breakdown { diff --git a/src/mcp/tools/common/history.c b/src/mcp/tools/common/history.c index 9f08e84..0e600d4 100644 --- a/src/mcp/tools/common/history.c +++ b/src/mcp/tools/common/history.c @@ -245,6 +245,25 @@ tiny_history_index_build(const tiny_store_t *store) return index; } +uint64_t +tiny_history_index_memory_estimate(const tiny_history_index_t *index) +{ + uint64_t total; + size_t position; + if (!index) + { + return 0; + } + total = sizeof(*index) + index->capacity * sizeof(*index->buckets) + + index->slot_capacity * sizeof(*index->slots); + for (position = 0; position < index->count; ++position) + { + total += strlen(index->buckets[position].session_id) + 1; + total += index->buckets[position].capacity * sizeof(*index->buckets[position].entries); + } + return total; +} + void tiny_history_index_free(tiny_history_index_t *index) { diff --git a/src/mcp/tools/common/history.h b/src/mcp/tools/common/history.h index 3b23027..ed7d11a 100644 --- a/src/mcp/tools/common/history.h +++ b/src/mcp/tools/common/history.h @@ -20,6 +20,7 @@ bool tiny_history_index_session(const tiny_history_index_t *index, const size_t **entries, size_t *count); size_t tiny_history_index_session_count(const tiny_history_index_t *index); +uint64_t tiny_history_index_memory_estimate(const tiny_history_index_t *index); bool tiny_history_index_session_at(const tiny_history_index_t *index, size_t position, const char **session_id, diff --git a/src/mcp/tools/common/memory_index.c b/src/mcp/tools/common/memory_index.c index e3841e2..32a97d1 100644 --- a/src/mcp/tools/common/memory_index.c +++ b/src/mcp/tools/common/memory_index.c @@ -165,6 +165,25 @@ tiny_memory_index_build(const tiny_store_t *store) return index; } +uint64_t +tiny_memory_index_memory_estimate(const tiny_memory_index_t *index) +{ + uint64_t total; + size_t position; + if (!index) + { + return 0; + } + total = sizeof(*index) + index->count * sizeof(*index->buckets) + + index->slot_capacity * sizeof(*index->slots); + for (position = 0; position < index->count; ++position) + { + total += index->buckets[position].incoming_capacity + * sizeof(*index->buckets[position].incoming); + } + return total; +} + void tiny_memory_index_free(tiny_memory_index_t *index) { diff --git a/src/mcp/tools/common/memory_index.h b/src/mcp/tools/common/memory_index.h index fbe43a7..71e450d 100644 --- a/src/mcp/tools/common/memory_index.h +++ b/src/mcp/tools/common/memory_index.h @@ -20,5 +20,6 @@ bool tiny_memory_index_incoming(const tiny_memory_index_t *index, const char *message_id, const tiny_relation_reference_t **references, size_t *count); +uint64_t tiny_memory_index_memory_estimate(const tiny_memory_index_t *index); #endif diff --git a/src/mcp/tools/get_system_health.c b/src/mcp/tools/get_system_health.c index f84511f..893a3b1 100644 --- a/src/mcp/tools/get_system_health.c +++ b/src/mcp/tools/get_system_health.c @@ -35,6 +35,25 @@ call_get_system_health(const mcpc_tool_t *tool, mcpc_ucbr_t *result) tiny_json_u64(&json, health.update_records); tiny_json_raw(&json, ",\"dead_bytes\":"); tiny_json_u64(&json, health.dead_bytes); + tiny_json_raw(&json, ",\"memory_estimate\":{\"total_bytes\":"); + tiny_json_u64(&json, + health.memory_entry_bytes + health.memory_body_bytes + health.memory_metadata_bytes + + tiny_search_index_memory_estimate(context->search_index) + + tiny_history_index_memory_estimate(context->history_index) + + tiny_memory_index_memory_estimate(context->memory_index)); + tiny_json_raw(&json, ",\"store_entries_bytes\":"); + tiny_json_u64(&json, health.memory_entry_bytes); + tiny_json_raw(&json, ",\"body_bytes\":"); + tiny_json_u64(&json, health.memory_body_bytes); + tiny_json_raw(&json, ",\"metadata_bytes\":"); + tiny_json_u64(&json, health.memory_metadata_bytes); + tiny_json_raw(&json, ",\"search_index_bytes\":"); + tiny_json_u64(&json, tiny_search_index_memory_estimate(context->search_index)); + tiny_json_raw(&json, ",\"history_index_bytes\":"); + tiny_json_u64(&json, tiny_history_index_memory_estimate(context->history_index)); + tiny_json_raw(&json, ",\"relation_index_bytes\":"); + tiny_json_u64(&json, tiny_memory_index_memory_estimate(context->memory_index)); + tiny_json_raw(&json, "}"); tiny_json_raw(&json, ",\"io\":{\"write_bytes\":"); tiny_json_u64(&json, health.write_bytes); tiny_json_raw(&json, ",\"write_calls\":"); diff --git a/src/storage/store.c b/src/storage/store.c index 060c6f6..e7ade28 100644 --- a/src/storage/store.c +++ b/src/storage/store.c @@ -1963,6 +1963,29 @@ tiny_store_health(const tiny_store_t *store, tiny_store_health_t *health) } for (index = 0; index < store->entry_count; ++index) { + const tiny_message_t *message = &store->entries[index].message; + size_t item; + health->memory_entry_bytes += sizeof(store->entries[index]); + health->memory_body_bytes += message->session_id ? strlen(message->session_id) + 1U : 0; + health->memory_body_bytes += message->content ? strlen(message->content) + 1U : 0; + health->memory_metadata_bytes + += message->metadata.tool_call_id ? strlen(message->metadata.tool_call_id) + 1U : 0; + health->memory_metadata_bytes + += message->metadata.summary ? strlen(message->metadata.summary) + 1U : 0; + health->memory_metadata_bytes + += message->metadata.relation_count * sizeof(*message->metadata.relations); + for (item = 0; item < message->metadata.tag_count; ++item) + { + health->memory_metadata_bytes += strlen(message->metadata.tags[item]) + 1U; + } + health->memory_metadata_bytes + += message->metadata.tag_count * sizeof(*message->metadata.tags); + health->memory_metadata_bytes + += message->metadata.source_count * sizeof(*message->metadata.source_message_ids); + for (item = 0; item < message->metadata.source_count; ++item) + { + health->memory_metadata_bytes += strlen(message->metadata.source_message_ids[item]) + 1U; + } if (store->entries[index].deleted) { health->deleted_messages++; diff --git a/src/storage/store.h b/src/storage/store.h index adeae35..9efb02a 100644 --- a/src/storage/store.h +++ b/src/storage/store.h @@ -45,6 +45,9 @@ typedef struct tiny_store_health uint64_t last_sequence; uint64_t update_records; uint64_t dead_bytes; + uint64_t memory_entry_bytes; + uint64_t memory_body_bytes; + uint64_t memory_metadata_bytes; uint64_t write_bytes; uint64_t write_calls; uint64_t flush_calls; -- cgit v1.2.3