Missing implementation in runAfterAgentCallbacks for invocation ended state

[MRuhan17]

Issue Description

In the runAfterAgentCallbacks function (agent/agent.go, line 264-279), there is a TODO comment indicating incomplete implementation:

// TODO set context invocation ended
// ctx.invocationEnded = true

This TODO appears at line 279 in the code:

func runAfterAgentCallbacks(ctx InvocationContext) (*session.Event, error) {
    // ... code omitted ...
    if newContent == nil {
        continue
    }
    event := session.NewEvent(ctx.InvocationID())
    event.LLMResponse = model.LLMResponse{
        Content: newContent,
    }
    event.Author = agent.Name()
    event.Branch = ctx.Branch()
    event.Actions = *callbackCtx.actions
    // TODO set context invocation ended      <--- HERE
    // ctx.invocationEnded = true
    return event, nil
}

Problem

The invocation ended state is NOT being set when an after-agent callback returns content. This could lead to:

1. Inconsistent state management between before and after callbacks
2. Potential for subsequent events to be processed incorrectly if the invocation should have ended
3. State tracking issues for session management

Expected Behavior

When an AfterAgentCallback returns non-nil content, the invocation should be marked as ended (similar to how it's done in runBeforeAgentCallbacks).

Location

File: agent/agent.go Line: 279

[MRuhan17]

This issue has been resolved in PR #371. I've implemented the missing ctx.EndInvocation() call in the runAfterAgentCallbacks() function to properly mark the invocation as ended when an after-agent callback returns content.

The fix ensures consistent state management with the before-agent callbacks and aligns with the proper invocation lifecycle.

Once PR #371 is merged, this issue will be fully resolved.

[baptmont]

Nice catch, looking at the python implementation I would say that could even be removed. There aren't ctx.Ended() validations after running the after agents callback and changes to InvocationContext are local so it doesn't affect execution on sequential agents.

[MRuhan17]

Thanks for the feedback! You raise a good point about comparing with the Python implementation. I understand your concern that since there aren't ctx.Ended() validations after the after-agents callback and changes to InvocationContext are local, the ctx.EndInvocation() call might seem unnecessary.

However, I'd like to understand your suggestion better:

1. Should we completely remove the call (matching Python behavior), or should we analyze if this is a difference in implementation approach between Python and Go?

2. State consistency: While the changes are local to the current context, should we still maintain consistency with how runBeforeAgentCallbacks() handles the invocation state for clarity and maintainability?

3. Sequential agents: You mention that it doesn't affect execution on sequential agents - could you elaborate on the execution flow to confirm this won't cause issues in edge cases?

I want to make sure we implement the most robust solution. If the Python implementation intentionally omits this, we should understand why before deciding if Go should do the same.

Let me know your thoughts and I can adjust the implementation accordingly!

[baptmont]

ctx.Ended is used to check if the invocation should continue or has ended. For example runBeforeAgentCallbacks calls EndInvocation if there is a content because it should return that event instead of running the agent. In the case of runAfterAgentCallback there is no need for it, because there are no validations after it that check the ended value. It does not affect sequential agents in the sense that ctx.Ended is restricted to the subagent and does not affect the parent, and this is the intended behavior.

[baptmont]

Matching python behavior is the correct approach.

[MRuhan17]

Thanks for the clarification. I’ve updated the PR to remove the EndInvocation() call and added a detailed comment explaining why the invocation state shouldn’t be modified in runAfterAgentCallbacks, matching the Python implementation and avoiding misleading behavior.

The PR now reflects the intended design. Let me know if anything else needs tightening.

[KirkifyAIGenerator]

thx