r/src/context_manager.c

// retoor <retoor@molodetz.nl>

#include "context_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MIN_KEEP_CHARS 500
#define TRUNCATE_MARKER "\n\n[... content truncated for context management ...]\n\n"

static const char *get_message_role(struct json_object *msg) {
    struct json_object *role_obj;
    if (json_object_object_get_ex(msg, "role", &role_obj)) {
        return json_object_get_string(role_obj);
    }
    return "";
}

static size_t get_message_content_len(struct json_object *msg) {
    struct json_object *content_obj;
    const char *content = NULL;

    if (json_object_object_get_ex(msg, "content", &content_obj)) {
        content = json_object_get_string(content_obj);
    } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
        content = json_object_get_string(content_obj);
    }

    return content ? strlen(content) : 0;
}

static size_t calculate_total_size(messages_handle msgs) {
    size_t total = 0;
    int count = messages_count(msgs);
    for (int i = 0; i < count; i++) {
        total += get_message_content_len(messages_get_object(msgs, i));
    }
    return total;
}

static r_status_t perform_truncate(messages_handle msgs, int index, double ratio) {
    struct json_object *msg = messages_get_object(msgs, index);
    if (!msg) return R_ERROR_NOT_FOUND;

    struct json_object *content_obj;
    const char *content = NULL;
    bool is_tool_result = false;

    if (json_object_object_get_ex(msg, "content", &content_obj)) {
        content = json_object_get_string(content_obj);
    } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
        content = json_object_get_string(content_obj);
        is_tool_result = true;
    }

    if (!content) return R_SUCCESS;

    size_t len = strlen(content);
    size_t target_len = (size_t)(len * ratio);
    if (target_len < MIN_KEEP_CHARS * 2) target_len = MIN_KEEP_CHARS * 2;
    if (target_len >= len) return R_SUCCESS;

    size_t keep_each = target_len / 2;
    
    char *new_content = malloc(keep_each * 2 + strlen(TRUNCATE_MARKER) + 1);
    if (!new_content) return R_ERROR_OUT_OF_MEMORY;

    strncpy(new_content, content, keep_each);
    new_content[keep_each] = '\0';
    strcat(new_content, TRUNCATE_MARKER);
    strcat(new_content, content + len - keep_each);

    struct json_object *new_msg = json_object_get(msg);
    if (is_tool_result) {
        json_object_object_add(new_msg, "tool_result", json_object_new_string(new_content));
    } else {
        json_object_object_add(new_msg, "content", json_object_new_string(new_content));
    }

    free(new_content);
    return messages_replace_at(msgs, index, new_msg);
}

r_status_t context_manager_shrink(messages_handle msgs) {
    if (!msgs) return R_ERROR_INVALID_ARG;

    int count = messages_count(msgs);
    if (count <= 1) return R_ERROR_API_ERROR;

    size_t initial_size = calculate_total_size(msgs);
    // Target 40% of initial size to be safe and avoid immediate re-overflow
    size_t target_size = (size_t)(initial_size * 0.4);
    if (target_size < 10000) target_size = 10000; // Don't shrink too much if it's already small

    fprintf(stderr, "  \033[2m-> Context overflow (approx %zu chars). Shrinking to %zu...\033[0m\n", 
            initial_size, target_size);

    int iterations = 0;
    while (calculate_total_size(msgs) > target_size && iterations < 50) {
        iterations++;
        count = messages_count(msgs);
        
        // Strategy 1: Find largest non-system, non-last message and truncate it
        int largest_idx = -1;
        size_t largest_size = 0;
        
        for (int i = 0; i < count - 1; i++) {
            struct json_object *msg = messages_get_object(msgs, i);
            if (strcmp(get_message_role(msg), "system") == 0) continue;
            
            size_t s = get_message_content_len(msg);
            if (s > largest_size) {
                largest_size = s;
                largest_idx = i;
            }
        }

        if (largest_idx != -1 && largest_size > 1000) {
            perform_truncate(msgs, largest_idx, 0.3); // Cut to 30% of its size
            continue;
        }

        // Strategy 2: Remove oldest removable messages (keep sequence)
        int first_removable = -1;
        for (int i = 0; i < count - 1; i++) {
            struct json_object *msg = messages_get_object(msgs, i);
            if (strcmp(get_message_role(msg), "system") != 0) {
                first_removable = i;
                break;
            }
        }

        if (first_removable != -1 && first_removable < count - 1) {
            // Remove 1 message at a time from the front
            messages_remove_range(msgs, first_removable, 1);
        } else {
            // Nothing left to remove but system or last message
            break;
        }
    }

    // Last Resort: If still too big, truncate the last message
    size_t final_size = calculate_total_size(msgs);
    if (final_size > target_size) {
        count = messages_count(msgs);
        if (count > 0) {
            perform_truncate(msgs, count - 1, 0.5);
        }
    }

    size_t shrunk_size = calculate_total_size(msgs);
    fprintf(stderr, "  \033[2m-> Context shrunk to approx %zu chars.\033[0m\n", shrunk_size);

    return R_SUCCESS;
}